diff --git a/.gitattributes b/.gitattributes
index 4fe091848b8698d7dc1250b949d7aa2e53bf5612..41cec71ce7e7c61ba4bb79dc149169e60d99b3de 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -76,3 +76,9 @@ me-seed-anneal_25000_6/ocdbt.process_0/d/4ff3129e369778ef3c5d569989c1b4e2 filter
 me-seed-anneal_25000_6/ocdbt.process_0/d/5ab60036321995fb12676c0c8c180126 filter=lfs diff=lfs merge=lfs -text
 me-seed-anneal_25000_6/ocdbt.process_0/d/9cc206d9b16ca8b229a5772172d855a0 filter=lfs diff=lfs merge=lfs -text
 me-seed-anneal_25000_6/ocdbt.process_0/d/d0e8edef8545dc8a0016de23f29d1c3a filter=lfs diff=lfs merge=lfs -text
+me-seed-anneal_25000_8/d/fc6532cee37a682ff2f0e598e35baf0a filter=lfs diff=lfs merge=lfs -text
+me-seed-anneal_25000_8/ocdbt.process_0/d/2141286350db8c5c8842f46870d0b713 filter=lfs diff=lfs merge=lfs -text
+me-seed-anneal_25000_8/ocdbt.process_0/d/66169ba653fa977241780f1de73151eb filter=lfs diff=lfs merge=lfs -text
+me-seed-anneal_25000_8/ocdbt.process_0/d/a242e476e90b40485e2d4e1a468efa4e filter=lfs diff=lfs merge=lfs -text
+me-seed-anneal_25000_8/ocdbt.process_0/d/ed2ae4eadbd5abaea8e9cc59055eed98 filter=lfs diff=lfs merge=lfs -text
+me-seed-anneal_25000_8/ocdbt.process_0/d/f4df26ebeeda30cb1cde15a19c867031 filter=lfs diff=lfs merge=lfs -text
diff --git a/me-seed-anneal_25000_8/_CHECKPOINT_METADATA b/me-seed-anneal_25000_8/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..ebb04835d7d38e0ab54ff44e5956fab9746efe8a
--- /dev/null
+++ b/me-seed-anneal_25000_8/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786888059091931869, "commit_timestamp_nsecs": 1786888067577094153, "custom_metadata": {}}
\ No newline at end of file
diff --git a/me-seed-anneal_25000_8/_METADATA b/me-seed-anneal_25000_8/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..34e12fdaa3372cfbfc072b292f31e35a9fd1d8b5
--- /dev/null
+++ b/me-seed-anneal_25000_8/_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, 1]}}, "('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": [3]}}, "('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, 1]}}, "('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": [2, 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": [2, 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": [128]}}, "('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, 128]}}, "('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": [128]}}, "('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, 128]}}, "('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": [2, 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": [2, 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": [64]}}, "('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": [64]}}, "('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": [2, 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": [64, 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": [64]}}, "('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": [2, 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": [2, 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": [64, 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": [2, 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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64]}}, "('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": [2, 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": [64, 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": [64]}}, "('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": [2, 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": [2, 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": [64, 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": [2, 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": [64, 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": [64]}}, "('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": [64, 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": [3]}}, "('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, 1]}}, "('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": [3]}}, "('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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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": [1, 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, 2]}}, "('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, 3]}}, "('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": [1, 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, 2]}}, "('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, 3]}}, "('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": [1, 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": [1]}}, "('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, 3]}}, "('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": [1]}}, "('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": [1, 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": [4]}}, "('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": [4, 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": [4, 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": [4, 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": [1, 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, 3]}}, "('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": [1, 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": [1]}}, "('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, 3]}}, "('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, 1]}}, "('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, 2]}}, "('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, 1, 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, 3]}}, "('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, 3]}}, "('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, 2]}}, "('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, 3]}}, "('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, 3]}}, "('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, 3]}}, "('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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [10240]}}, "('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": [3]}}, "('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": [20480]}}, "('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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [2, 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": [64, 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": [64]}}, "('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": [2, 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": [2, 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": [64, 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": [2, 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": [64, 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": [2, 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": [64, 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": [64]}}, "('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": [2, 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": [2, 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": [64, 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": [2, 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": [64, 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, 1]}}, "('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": [3]}}, "('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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [1, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [10240]}}, "('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": [3]}}, "('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": [20480]}}, "('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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [2, 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": [64, 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": [64]}}, "('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": [2, 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": [2, 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": [64, 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": [2, 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": [64, 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": [2, 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": [64, 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": [64]}}, "('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": [2, 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": [2, 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": [64, 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": [2, 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": [64, 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, 1]}}, "('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": [3]}}, "('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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [1, 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": [3]}}, "('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, 3, 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, 3, 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": [3]}}, "('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": [1, 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": [3]}}, "('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": [1, 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": [3]}}, "('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, 3, 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, 3, 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": [3]}}, "('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": [1, 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": [3]}}, "('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": [1, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [64, 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": [3, 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": [128]}}, "('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, 128]}}, "('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": [32]}}, "('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": [32, 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": [3, 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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [3]}}, "('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, 3, 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, 3, 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": [3]}}, "('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": [1, 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": [3]}}, "('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": [1, 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": [3]}}, "('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, 3, 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, 3, 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": [3]}}, "('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": [1, 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": [3]}}, "('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": [1, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [64, 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": [3, 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": [128]}}, "('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, 128]}}, "('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": [32]}}, "('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": [32, 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": [3, 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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [3]}}, "('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, 3, 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, 3, 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": [3]}}, "('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": [1, 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": [3]}}, "('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": [1, 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": [3]}}, "('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, 3, 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, 3, 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": [3]}}, "('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": [1, 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": [3]}}, "('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": [1, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [128, 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": [3, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [3, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [3]}}, "('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, 3, 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, 3, 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": [3]}}, "('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": [1, 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": [3]}}, "('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": [1, 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": [3]}}, "('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, 3, 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, 3, 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": [3]}}, "('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": [1, 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": [3]}}, "('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": [1, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [128, 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": [3, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [3, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [3]}}, "('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, 3, 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, 3, 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": [3]}}, "('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": [1, 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": [3]}}, "('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": [1, 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": [3]}}, "('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, 3, 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, 3, 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": [3]}}, "('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": [1, 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": [3]}}, "('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": [1, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [128, 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": [3, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [3, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [3]}}, "('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, 3, 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, 3, 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": [3]}}, "('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": [1, 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": [3]}}, "('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": [1, 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": [3]}}, "('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, 3, 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, 3, 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": [3]}}, "('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": [1, 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": [3]}}, "('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": [1, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [128, 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": [3, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [3, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [3]}}, "('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, 3, 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, 3, 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": [3]}}, "('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": [1, 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": [3]}}, "('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": [1, 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": [3]}}, "('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, 3, 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, 3, 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": [3]}}, "('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": [1, 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": [3]}}, "('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": [1, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [128, 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": [3, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [3, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [3]}}, "('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, 3, 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, 3, 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": [3]}}, "('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": [1, 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": [3]}}, "('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": [1, 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": [3]}}, "('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, 3, 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, 3, 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": [3]}}, "('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": [1, 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": [3]}}, "('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": [1, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [128, 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": [3, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [3, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [3]}}, "('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, 3, 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, 3, 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": [3]}}, "('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": [1, 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": [3]}}, "('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": [1, 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": [3]}}, "('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, 3, 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, 3, 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": [3]}}, "('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": [1, 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": [3]}}, "('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": [1, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [64, 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": [3, 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": [128]}}, "('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, 128]}}, "('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": [32]}}, "('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": [32, 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": [3, 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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [3]}}, "('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": [1, 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": [3]}}, "('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, 3, 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, 3, 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": [3]}}, "('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": [1, 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": [3]}}, "('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": [1, 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": [3]}}, "('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, 3, 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, 3, 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": [3]}}, "('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": [1, 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": [3]}}, "('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": [1, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [128, 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": [3, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [3, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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, 1]}}, "('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, 3]}}, "('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": [3]}}, "('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": [3]}}, "('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": [1, 64, 64]}}, "('0', 'unet', 'seed_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "seed_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', 'seed_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "seed_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', 'seed_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "seed_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', 'seed_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "seed_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, 32]}}, "('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": [1, 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, 1]}}, "('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, 1]}}, "('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": [128]}}, "('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, 128]}}, "('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": [128]}}, "('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, 128]}}, "('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": [64]}}, "('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": [64]}}, "('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": [2, 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": [64, 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": [64]}}, "('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": [2, 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": [2, 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": [64, 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": [2, 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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64]}}, "('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": [2, 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": [64, 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": [64]}}, "('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": [2, 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": [2, 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": [64, 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": [2, 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": [64, 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": [64]}}, "('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": [64, 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, 1]}}, "('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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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": [1, 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, 3]}}, "('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": [1, 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, 3]}}, "('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, 3]}}, "('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, 3]}}, "('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, 3]}}, "('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, 3]}}, "('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, 3]}}, "('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, 3]}}, "('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, 3]}}, "('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, 3]}}, "('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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [10240]}}, "('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": [20480]}}, "('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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [2, 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": [64, 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": [64]}}, "('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": [2, 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": [2, 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": [64, 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": [2, 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": [64, 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": [2, 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": [64, 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": [64]}}, "('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": [2, 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": [2, 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": [64, 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": [2, 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": [64, 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, 1]}}, "('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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [10240]}}, "('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": [20480]}}, "('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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [2, 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": [64, 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": [64]}}, "('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": [2, 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": [2, 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": [64, 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": [2, 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": [64, 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": [2, 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": [64, 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": [64]}}, "('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": [2, 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": [2, 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": [64, 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": [2, 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": [64, 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, 1]}}, "('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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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": [1, 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": [1, 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": [1, 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, 1]}}, "('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": [64, 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": [128]}}, "('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, 128]}}, "('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": [32]}}, "('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": [32, 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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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": [1, 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": [1, 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": [1, 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, 1]}}, "('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": [64, 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": [128]}}, "('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, 128]}}, "('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": [32]}}, "('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": [32, 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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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": [1, 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": [1, 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": [1, 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, 1]}}, "('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": [128, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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": [1, 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": [1, 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": [1, 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, 1]}}, "('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": [128, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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": [1, 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": [1, 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": [1, 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, 1]}}, "('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": [128, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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": [1, 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": [1, 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": [1, 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, 1]}}, "('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": [128, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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": [1, 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": [1, 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": [1, 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, 1]}}, "('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": [128, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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": [1, 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": [1, 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": [1, 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, 1]}}, "('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": [128, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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": [1, 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": [1, 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": [1, 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, 1]}}, "('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": [64, 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": [128]}}, "('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, 128]}}, "('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": [32]}}, "('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": [32, 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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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": [1, 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": [1, 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": [1, 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": [1, 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, 1]}}, "('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": [128, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'seed_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": "seed_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', 'seed_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": "seed_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', 'seed_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": "seed_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', 'seed_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": "seed_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, 32]}}, "('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, 1]}}, "('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, 1]}}, "('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": [128]}}, "('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, 128]}}, "('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": [128]}}, "('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, 128]}}, "('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": [64]}}, "('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": [64]}}, "('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": [2, 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": [64, 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": [64]}}, "('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": [2, 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": [2, 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": [64, 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": [2, 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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64]}}, "('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": [2, 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": [64, 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": [64]}}, "('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": [2, 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": [2, 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": [64, 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": [2, 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": [64, 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": [64]}}, "('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": [64, 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, 1]}}, "('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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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": [1, 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, 3]}}, "('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": [1, 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, 3]}}, "('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, 3]}}, "('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, 3]}}, "('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, 3]}}, "('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, 3]}}, "('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, 3]}}, "('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, 3]}}, "('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, 3]}}, "('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, 3]}}, "('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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [10240]}}, "('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": [20480]}}, "('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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [2, 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": [64, 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": [64]}}, "('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": [2, 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": [2, 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": [64, 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": [2, 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": [64, 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": [2, 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": [64, 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": [64]}}, "('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": [2, 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": [2, 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": [64, 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": [2, 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": [64, 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, 1]}}, "('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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [10240]}}, "('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": [20480]}}, "('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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [2, 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": [64, 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": [64]}}, "('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": [2, 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": [2, 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": [64, 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": [2, 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": [64, 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": [2, 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": [64, 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": [64]}}, "('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": [2, 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": [2, 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": [64, 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": [2, 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": [64, 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, 1]}}, "('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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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": [1, 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": [1, 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": [1, 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, 1]}}, "('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": [64, 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": [128]}}, "('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, 128]}}, "('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": [32]}}, "('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": [32, 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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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": [1, 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": [1, 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": [1, 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, 1]}}, "('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": [64, 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": [128]}}, "('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, 128]}}, "('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": [32]}}, "('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": [32, 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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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": [1, 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": [1, 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": [1, 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, 1]}}, "('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": [128, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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": [1, 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": [1, 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": [1, 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, 1]}}, "('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": [128, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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": [1, 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": [1, 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": [1, 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, 1]}}, "('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": [128, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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": [1, 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": [1, 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": [1, 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, 1]}}, "('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": [128, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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": [1, 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": [1, 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": [1, 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, 1]}}, "('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": [128, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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": [1, 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": [1, 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": [1, 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, 1]}}, "('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": [128, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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": [1, 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": [1, 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": [1, 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, 1]}}, "('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": [64, 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": [128]}}, "('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, 128]}}, "('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": [32]}}, "('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": [32, 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": [64]}}, "('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, 64]}}, "('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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 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": [1, 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": [1, 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": [1, 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": [1, 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, 1]}}, "('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": [128, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [256]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [64]}}, "('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": [64, 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": [128]}}, "('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": [64, 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, 1]}}, "('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, 1]}}, "('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, 1]}}, "('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": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'seed_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": "seed_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', 'seed_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": "seed_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', 'seed_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": "seed_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', 'seed_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": "seed_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, 32]}}, "('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-seed-anneal_25000_8/_sharding b/me-seed-anneal_25000_8/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..a41c54a72b12bc2f8c6f57135c83c260b8a7b589
--- /dev/null
+++ b/me-seed-anneal_25000_8/_sharding
@@ -0,0 +1 @@
+{"MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LmF0dG5fd2VpZ2h0cw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAucHJvai5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEucHJvai5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnNlZWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnNlZWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnNlZWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnNlZWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LndlaWdodF9pcnJlcC5ydG9p":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmxuX2Rpc3RfZG9mLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmxuX2Rpc3RfZG9mLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmxuX3RtcF9kb2YuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmxuX3RtcF9kb2Yuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmxuX3RtcF9wb3NlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmxuX3RtcF9wb3NlLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuc2RfZG90X21vZHVsYXRpb24=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuYWxpZ25fdG9fYXBwcm9hY2guMA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuYWxpZ25fdG9fYXBwcm9hY2guMQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuZmluZ2VydGlwX25vcm1hbHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuZmluZ2VydGlwX2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuaW5pdF9wcmVncmFzcF9qb2ludA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIubG9jYWxfZmluZ2VydGlwX2NvbnRhY3RfcG9zaXRpb25z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuYWxpZ25fdG9fYXBwcm9hY2guMA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuYWxpZ25fdG9fYXBwcm9hY2guMQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuZmluZ2VydGlwX25vcm1hbHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuZmluZ2VydGlwX2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuaW5pdF9wcmVncmFzcF9qb2ludA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMubG9jYWxfZmluZ2VydGlwX2NvbnRhY3RfcG9zaXRpb25z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3Rfc2hfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcud2VpZ2h0X2lycmVwLnJ0b2k=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi53ZWlnaHRfaXJyZXAucnRvaQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC53ZWlnaHRfaXJyZXAucnRvaQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi53ZWlnaHRfaXJyZXAucnRvaQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC53ZWlnaHRfaXJyZXAucnRvaQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5jb3VudA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LmF0dG5fd2VpZ2h0cw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnNlZWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnNlZWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnNlZWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnNlZWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX2Rpc3RfZG9mLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX2Rpc3RfZG9mLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX3RtcF9kb2YuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX3RtcF9kb2Yuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX3RtcF9wb3NlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX3RtcF9wb3NlLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuc2RfZG90X21vZHVsYXRpb24=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3Rfc2hfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LmF0dG5fd2VpZ2h0cw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnNlZWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnNlZWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnNlZWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnNlZWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX2Rpc3RfZG9mLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX2Rpc3RfZG9mLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX3RtcF9kb2YuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX3RtcF9kb2Yuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX3RtcF9wb3NlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX3RtcF9wb3NlLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuc2RfZG90X21vZHVsYXRpb24=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3Rfc2hfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMi5jb3VudA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5zdGVw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}"}
\ No newline at end of file
diff --git a/me-seed-anneal_25000_8/array_metadatas/process_0 b/me-seed-anneal_25000_8/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..cab9b2c86ee1e3c5eddae4d89c33ae2629719b66
--- /dev/null
+++ b/me-seed-anneal_25000_8/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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "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": [2, 7, 9], "chunk_shape": [2, 7, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid.to_grid_mat", "write_shape": [2, 7, 9], "chunk_shape": [2, 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": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.fc_0.kernel", "write_shape": [256, 128], "chunk_shape": [256, 128], "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": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.4.kernel", "write_shape": [256, 128], "chunk_shape": [256, 128], "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": [2, 7, 9], "chunk_shape": [2, 7, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid.to_grid_mat", "write_shape": [2, 7, 9], "chunk_shape": [2, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.2.layers.1.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.2.layers.1.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.expand_index", "write_shape": [3], "chunk_shape": [3], "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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "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": [1, 64, 9], "chunk_shape": [1, 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, 2], "chunk_shape": [2, 2], "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, 3], "chunk_shape": [2, 64, 3], "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": [1, 64, 9], "chunk_shape": [1, 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, 2], "chunk_shape": [2, 2], "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, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.align_to_approach.0", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.align_to_approach.1", "write_shape": [1], "chunk_shape": [1], "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, 3], "chunk_shape": [13, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.fingertip_idx", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.fingertip_normals", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.init_pregrasp_joint", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.joint_ranges", "write_shape": [4, 2], "chunk_shape": [4, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.joint_transforms", "write_shape": [4, 6], "chunk_shape": [4, 6], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.kinematics_transforms", "write_shape": [4, 7], "chunk_shape": [4, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.local_fingertip_contact_positions", "write_shape": [1, 6, 3], "chunk_shape": [1, 6, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.align_to_approach.0", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.align_to_approach.1", "write_shape": [1], "chunk_shape": [1], "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, 3], "chunk_shape": [17, 64, 3], "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, 1], "chunk_shape": [4, 1], "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, 2], "chunk_shape": [16, 2], "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, 1, 3], "chunk_shape": [4, 1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "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, 3], "chunk_shape": [23, 64, 3], "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, 2], "chunk_shape": [22, 2], "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, 3], "chunk_shape": [2, 64, 3], "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, 3], "chunk_shape": [1, 64, 3], "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, 3], "chunk_shape": [2, 64, 3], "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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.q_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.q_dot_sh_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.weight_irrep.rtoi", "write_shape": [1, 9], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.0.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.0.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [32, 1], "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, 3], "chunk_shape": [1, 3], "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": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "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, 3], "chunk_shape": [1, 3], "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": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "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, 3], "chunk_shape": [1, 3], "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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.s_proj.l_weights", "write_shape": [10240], "chunk_shape": [10240], "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": [3], "chunk_shape": [3], "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": [20480], "chunk_shape": [20480], "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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.q_dot_sh_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.q_s_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.weight_irrep.rtoi", "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.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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_feat.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_feat.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_query.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_query.expand_index", "write_shape": [3], "chunk_shape": [3], "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": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 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": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 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": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 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": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 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": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 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": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.expand_index", "write_shape": [3], "chunk_shape": [3], "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": [64], "chunk_shape": [64], "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, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "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": [1, 9], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.0.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.0.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [32, 1], "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, 3], "chunk_shape": [1, 3], "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": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "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, 3], "chunk_shape": [1, 3], "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": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "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, 3], "chunk_shape": [1, 3], "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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.s_proj.l_weights", "write_shape": [10240], "chunk_shape": [10240], "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": [3], "chunk_shape": [3], "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": [20480], "chunk_shape": [20480], "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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.q_dot_sh_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.q_s_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.weight_irrep.rtoi", "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.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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_feat.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_feat.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_query.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_query.expand_index", "write_shape": [3], "chunk_shape": [3], "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": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 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": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 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": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 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": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 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": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 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": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.expand_index", "write_shape": [3], "chunk_shape": [3], "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": [64], "chunk_shape": [64], "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, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "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": [1, 9], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 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, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "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": [64, 224], "chunk_shape": [64, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 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": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.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.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": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.kernel", "write_shape": [32, 96], "chunk_shape": [32, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 64], "chunk_shape": [128, 64], "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, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.tp_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.tp_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_ffn.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_ffn.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 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, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "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": [64, 224], "chunk_shape": [64, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 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": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.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.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": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.kernel", "write_shape": [32, 96], "chunk_shape": [32, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 64], "chunk_shape": [128, 64], "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, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.tp_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.tp_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_ffn.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_ffn.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "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": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.tp_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.tp_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_ffn.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_ffn.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "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": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.tp_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.tp_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_ffn.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_ffn.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "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": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.tp_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.tp_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_ffn.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_ffn.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "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": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.tp_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.tp_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_ffn.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_ffn.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "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": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.tp_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.tp_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_ffn.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_ffn.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "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": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.tp_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.tp_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_ffn.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_ffn.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 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, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "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": [64, 224], "chunk_shape": [64, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 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": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.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.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": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.fc_0.kernel", "write_shape": [32, 96], "chunk_shape": [32, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 64], "chunk_shape": [128, 64], "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, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.tp_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.tp_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_ffn.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_ffn.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.proj.weight", "write_shape": [1, 64, 32], "chunk_shape": [1, 64, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "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": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.tp_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.tp_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_ffn.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_ffn.expand_index", "write_shape": [3], "chunk_shape": [3], "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, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "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": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.proj.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.seed_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.seed_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.seed_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.seed_mlp.layers.2.kernel", "write_shape": [64, 32], "chunk_shape": [64, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_irrep.rtoi", "write_shape": [1, 9], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.fc_0.kernel", "write_shape": [256, 128], "chunk_shape": [256, 128], "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": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.4.kernel", "write_shape": [256, 128], "chunk_shape": [256, 128], "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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.2.layers.1.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.2.layers.1.kernel", "write_shape": [64, 192], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.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.kinematic_encoding.gl.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.0.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.0.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.1.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.1.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.2.dof_embedding", "write_shape": [13, 64, 3], "chunk_shape": [13, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.2.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.3.dof_embedding", "write_shape": [17, 64, 3], "chunk_shape": [17, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.3.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.4.dof_embedding", "write_shape": [23, 64, 3], "chunk_shape": [23, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.4.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.5.dof_embedding", "write_shape": [1, 64, 3], "chunk_shape": [1, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.5.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.q_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.q_dot_sh_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [32, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [10240], "chunk_shape": [10240], "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": [20480], "chunk_shape": [20480], "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": [64], "chunk_shape": [64], "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, 64], "chunk_shape": [256, 64], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [2, 32, 192], "chunk_shape": [2, 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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [2, 32, 192], "chunk_shape": [2, 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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "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": [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.6.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.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [32, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [10240], "chunk_shape": [10240], "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": [20480], "chunk_shape": [20480], "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": [64], "chunk_shape": [64], "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, 64], "chunk_shape": [256, 64], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [2, 32, 192], "chunk_shape": [2, 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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [2, 32, 192], "chunk_shape": [2, 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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "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": [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.6.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.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [1, 32, 32], "chunk_shape": [1, 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": [1, 32, 32], "chunk_shape": [1, 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": [1, 32, 32], "chunk_shape": [1, 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": [1, 32, 32], "chunk_shape": [1, 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, 1], "chunk_shape": [32, 1], "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": [64, 224], "chunk_shape": [64, 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": [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.4.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.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": [32], "chunk_shape": [32], "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": [32, 96], "chunk_shape": [32, 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": [64], "chunk_shape": [64], "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, 64], "chunk_shape": [128, 64], "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, 1], "chunk_shape": [32, 1], "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, 1], "chunk_shape": [32, 1], "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, 1], "chunk_shape": [32, 1], "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": [1, 32, 32], "chunk_shape": [1, 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": [1, 32, 32], "chunk_shape": [1, 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": [1, 32, 32], "chunk_shape": [1, 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": [1, 32, 32], "chunk_shape": [1, 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, 1], "chunk_shape": [32, 1], "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": [64, 224], "chunk_shape": [64, 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": [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.4.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.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": [32], "chunk_shape": [32], "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": [32, 96], "chunk_shape": [32, 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": [64], "chunk_shape": [64], "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, 64], "chunk_shape": [128, 64], "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, 1], "chunk_shape": [32, 1], "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, 1], "chunk_shape": [32, 1], "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, 1], "chunk_shape": [32, 1], "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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "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": [128, 448], "chunk_shape": [128, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [256], "chunk_shape": [256], "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": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [128], "chunk_shape": [128], "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": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "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": [128, 448], "chunk_shape": [128, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [256], "chunk_shape": [256], "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": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [128], "chunk_shape": [128], "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": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "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": [128, 448], "chunk_shape": [128, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [256], "chunk_shape": [256], "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": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [128], "chunk_shape": [128], "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": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "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": [128, 448], "chunk_shape": [128, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [256], "chunk_shape": [256], "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": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [128], "chunk_shape": [128], "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": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "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": [128, 448], "chunk_shape": [128, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [256], "chunk_shape": [256], "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": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [128], "chunk_shape": [128], "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": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "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": [128, 448], "chunk_shape": [128, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [256], "chunk_shape": [256], "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": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [128], "chunk_shape": [128], "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": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [1, 32, 32], "chunk_shape": [1, 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": [1, 32, 32], "chunk_shape": [1, 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": [1, 32, 32], "chunk_shape": [1, 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": [1, 32, 32], "chunk_shape": [1, 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, 1], "chunk_shape": [32, 1], "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": [64, 224], "chunk_shape": [64, 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": [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.4.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.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": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.fc_0.kernel", "write_shape": [32, 96], "chunk_shape": [32, 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": [64], "chunk_shape": [64], "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, 64], "chunk_shape": [128, 64], "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, 1], "chunk_shape": [32, 1], "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, 1], "chunk_shape": [32, 1], "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, 1], "chunk_shape": [32, 1], "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": [1, 64, 32], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "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": [128, 448], "chunk_shape": [128, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [256], "chunk_shape": [256], "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": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [128], "chunk_shape": [128], "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": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.seed_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.seed_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.seed_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.seed_mlp.layers.2.kernel", "write_shape": [64, 32], "chunk_shape": [64, 32], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.fc_0.kernel", "write_shape": [256, 128], "chunk_shape": [256, 128], "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": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.4.kernel", "write_shape": [256, 128], "chunk_shape": [256, 128], "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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.2.layers.1.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.2.layers.1.kernel", "write_shape": [64, 192], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.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.kinematic_encoding.gl.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.0.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.0.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.1.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.1.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.2.dof_embedding", "write_shape": [13, 64, 3], "chunk_shape": [13, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.2.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.3.dof_embedding", "write_shape": [17, 64, 3], "chunk_shape": [17, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.3.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.4.dof_embedding", "write_shape": [23, 64, 3], "chunk_shape": [23, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.4.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.5.dof_embedding", "write_shape": [1, 64, 3], "chunk_shape": [1, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.5.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.q_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.q_dot_sh_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [32, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [10240], "chunk_shape": [10240], "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": [20480], "chunk_shape": [20480], "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": [64], "chunk_shape": [64], "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, 64], "chunk_shape": [256, 64], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [2, 32, 192], "chunk_shape": [2, 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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [2, 32, 192], "chunk_shape": [2, 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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "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": [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.6.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.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [32, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [10240], "chunk_shape": [10240], "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": [20480], "chunk_shape": [20480], "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": [64], "chunk_shape": [64], "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, 64], "chunk_shape": [256, 64], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [2, 32, 192], "chunk_shape": [2, 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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [2, 32, 192], "chunk_shape": [2, 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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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": [2, 32], "chunk_shape": [2, 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": [64, 6, 32], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "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": [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.6.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.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [1, 32, 32], "chunk_shape": [1, 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": [1, 32, 32], "chunk_shape": [1, 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": [1, 32, 32], "chunk_shape": [1, 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": [1, 32, 32], "chunk_shape": [1, 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, 1], "chunk_shape": [32, 1], "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": [64, 224], "chunk_shape": [64, 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": [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.4.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.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": [32], "chunk_shape": [32], "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": [32, 96], "chunk_shape": [32, 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": [64], "chunk_shape": [64], "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, 64], "chunk_shape": [128, 64], "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, 1], "chunk_shape": [32, 1], "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, 1], "chunk_shape": [32, 1], "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, 1], "chunk_shape": [32, 1], "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": [1, 32, 32], "chunk_shape": [1, 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": [1, 32, 32], "chunk_shape": [1, 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": [1, 32, 32], "chunk_shape": [1, 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": [1, 32, 32], "chunk_shape": [1, 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, 1], "chunk_shape": [32, 1], "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": [64, 224], "chunk_shape": [64, 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": [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.4.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.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": [32], "chunk_shape": [32], "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": [32, 96], "chunk_shape": [32, 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": [64], "chunk_shape": [64], "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, 64], "chunk_shape": [128, 64], "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, 1], "chunk_shape": [32, 1], "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, 1], "chunk_shape": [32, 1], "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, 1], "chunk_shape": [32, 1], "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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "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": [128, 448], "chunk_shape": [128, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [256], "chunk_shape": [256], "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": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [128], "chunk_shape": [128], "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": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "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": [128, 448], "chunk_shape": [128, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [256], "chunk_shape": [256], "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": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [128], "chunk_shape": [128], "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": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "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": [128, 448], "chunk_shape": [128, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [256], "chunk_shape": [256], "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": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [128], "chunk_shape": [128], "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": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "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": [128, 448], "chunk_shape": [128, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [256], "chunk_shape": [256], "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": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [128], "chunk_shape": [128], "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": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "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": [128, 448], "chunk_shape": [128, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [256], "chunk_shape": [256], "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": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [128], "chunk_shape": [128], "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": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "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": [128, 448], "chunk_shape": [128, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [256], "chunk_shape": [256], "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": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [128], "chunk_shape": [128], "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": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [1, 32, 32], "chunk_shape": [1, 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": [1, 32, 32], "chunk_shape": [1, 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": [1, 32, 32], "chunk_shape": [1, 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": [1, 32, 32], "chunk_shape": [1, 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, 1], "chunk_shape": [32, 1], "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": [64, 224], "chunk_shape": [64, 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": [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.4.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.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": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.fc_0.kernel", "write_shape": [32, 96], "chunk_shape": [32, 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": [64], "chunk_shape": [64], "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, 64], "chunk_shape": [128, 64], "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, 1], "chunk_shape": [32, 1], "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, 1], "chunk_shape": [32, 1], "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, 1], "chunk_shape": [32, 1], "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": [1, 64, 32], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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": [1, 64, 64], "chunk_shape": [1, 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, 1], "chunk_shape": [64, 1], "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": [128, 448], "chunk_shape": [128, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [256], "chunk_shape": [256], "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": [64, 768], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [64], "chunk_shape": [64], "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": [64, 192], "chunk_shape": [64, 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": [128], "chunk_shape": [128], "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": [64, 384], "chunk_shape": [64, 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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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, 1], "chunk_shape": [64, 1], "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": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.seed_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.seed_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.seed_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.seed_mlp.layers.2.kernel", "write_shape": [64, 32], "chunk_shape": [64, 32], "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-seed-anneal_25000_8/d/fc6532cee37a682ff2f0e598e35baf0a b/me-seed-anneal_25000_8/d/fc6532cee37a682ff2f0e598e35baf0a
new file mode 100644
index 0000000000000000000000000000000000000000..089d5a12461bec23dbb94be4b8bff8868fda757e
--- /dev/null
+++ b/me-seed-anneal_25000_8/d/fc6532cee37a682ff2f0e598e35baf0a
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:11e130593e7089bc9c12a349494649d584fc44b8f5b1cb5931f5886128d3eecf
+size 840918
diff --git a/me-seed-anneal_25000_8/manifest.ocdbt b/me-seed-anneal_25000_8/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..5cf99c55171eafccdfe3a47e67afc09258f22de7
Binary files /dev/null and b/me-seed-anneal_25000_8/manifest.ocdbt differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/024f242a11354e2b60fc152abfe6dd4a b/me-seed-anneal_25000_8/ocdbt.process_0/d/024f242a11354e2b60fc152abfe6dd4a
new file mode 100644
index 0000000000000000000000000000000000000000..7d3e8c3c7a6f39e4b5b84fee75e6ace44f000a5b
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/024f242a11354e2b60fc152abfe6dd4a differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/035514618dc1027020e7c58110af0fe9 b/me-seed-anneal_25000_8/ocdbt.process_0/d/035514618dc1027020e7c58110af0fe9
new file mode 100644
index 0000000000000000000000000000000000000000..becef566ce64b5125731bd8e1b659ac4127e6659
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/035514618dc1027020e7c58110af0fe9 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/0444c3cbd549472ce9c95d26e62f6cf9 b/me-seed-anneal_25000_8/ocdbt.process_0/d/0444c3cbd549472ce9c95d26e62f6cf9
new file mode 100644
index 0000000000000000000000000000000000000000..4b0a8c53781d27903e5548d9104bbe31d03abe93
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/0444c3cbd549472ce9c95d26e62f6cf9 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/0b4c3578dd06da986de9202f5be46933 b/me-seed-anneal_25000_8/ocdbt.process_0/d/0b4c3578dd06da986de9202f5be46933
new file mode 100644
index 0000000000000000000000000000000000000000..af471862ea5fa74080d57f1c5babe9c8b2c5737b
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/0b4c3578dd06da986de9202f5be46933 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/11c4fce6e88e42db7d6a15c00cdecb7f b/me-seed-anneal_25000_8/ocdbt.process_0/d/11c4fce6e88e42db7d6a15c00cdecb7f
new file mode 100644
index 0000000000000000000000000000000000000000..6c9d9402bfe5be06cc1bbb94c4b8a78361258577
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/11c4fce6e88e42db7d6a15c00cdecb7f differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/12203eb2d13776ad8b39ec45d3ea78d6 b/me-seed-anneal_25000_8/ocdbt.process_0/d/12203eb2d13776ad8b39ec45d3ea78d6
new file mode 100644
index 0000000000000000000000000000000000000000..9b6ead5c5b8ec0fc68404ef0bfd17cc4cad30fb4
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/12203eb2d13776ad8b39ec45d3ea78d6 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/147634f311833edb9880e8cd5cdac0ca b/me-seed-anneal_25000_8/ocdbt.process_0/d/147634f311833edb9880e8cd5cdac0ca
new file mode 100644
index 0000000000000000000000000000000000000000..22057a9ba47adc82a90c189e68d1a728c0004d68
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/147634f311833edb9880e8cd5cdac0ca differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/158b862f408ee6e7f6612156245a50dc b/me-seed-anneal_25000_8/ocdbt.process_0/d/158b862f408ee6e7f6612156245a50dc
new file mode 100644
index 0000000000000000000000000000000000000000..b4859bae398ff4d365ebbb404289ccb6afd9ae06
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/158b862f408ee6e7f6612156245a50dc differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/16b98053c181da8ac83cc588697a69d6 b/me-seed-anneal_25000_8/ocdbt.process_0/d/16b98053c181da8ac83cc588697a69d6
new file mode 100644
index 0000000000000000000000000000000000000000..6f7cdc3d84978ee4b32b347245d14dfbfd442995
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/16b98053c181da8ac83cc588697a69d6 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/1f726a36fa850b269baadebb2a9dd2ab b/me-seed-anneal_25000_8/ocdbt.process_0/d/1f726a36fa850b269baadebb2a9dd2ab
new file mode 100644
index 0000000000000000000000000000000000000000..74f0e91868f1d24d26e9f9f3668170b21e3ac026
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/1f726a36fa850b269baadebb2a9dd2ab differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/2141286350db8c5c8842f46870d0b713 b/me-seed-anneal_25000_8/ocdbt.process_0/d/2141286350db8c5c8842f46870d0b713
new file mode 100644
index 0000000000000000000000000000000000000000..122fc10b83211e71c9f8fe82151db33a1b003310
--- /dev/null
+++ b/me-seed-anneal_25000_8/ocdbt.process_0/d/2141286350db8c5c8842f46870d0b713
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a6df9a5087011792fce7ecea6d75155d947dc7e08bcb62d536a0f8ef38de7567
+size 147636224
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/2257f06fdb29ff1d35f1c15588aacb20 b/me-seed-anneal_25000_8/ocdbt.process_0/d/2257f06fdb29ff1d35f1c15588aacb20
new file mode 100644
index 0000000000000000000000000000000000000000..28cab278342887ba82b4567a18804d745e3e24f8
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/2257f06fdb29ff1d35f1c15588aacb20 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/22c36f2fd594c6c7383e9952562e5039 b/me-seed-anneal_25000_8/ocdbt.process_0/d/22c36f2fd594c6c7383e9952562e5039
new file mode 100644
index 0000000000000000000000000000000000000000..076ba44d7e5d2c3247c284e7fd9cdbd4955c197b
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/22c36f2fd594c6c7383e9952562e5039 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/2f8b3a50cdec0019ae4fd67ef69768aa b/me-seed-anneal_25000_8/ocdbt.process_0/d/2f8b3a50cdec0019ae4fd67ef69768aa
new file mode 100644
index 0000000000000000000000000000000000000000..ebeba878fdfb5e5ee2fc955e891c017b325236d0
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/2f8b3a50cdec0019ae4fd67ef69768aa differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/32bf0494f25c7fe9dcfe955124f7da95 b/me-seed-anneal_25000_8/ocdbt.process_0/d/32bf0494f25c7fe9dcfe955124f7da95
new file mode 100644
index 0000000000000000000000000000000000000000..1ce71f9f7ddb881e133ba1b25c34e95f31f193af
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/32bf0494f25c7fe9dcfe955124f7da95 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/33cfbe7918ee294f5478b9fb40613daf b/me-seed-anneal_25000_8/ocdbt.process_0/d/33cfbe7918ee294f5478b9fb40613daf
new file mode 100644
index 0000000000000000000000000000000000000000..f365fbed6ab3e64086aec66d680c4045c5801330
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/33cfbe7918ee294f5478b9fb40613daf differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/344c4c92d3e6c174084e4ce2056f3d60 b/me-seed-anneal_25000_8/ocdbt.process_0/d/344c4c92d3e6c174084e4ce2056f3d60
new file mode 100644
index 0000000000000000000000000000000000000000..09d7536c2a04da023d25dc56a4b5dbb87c4471b7
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/344c4c92d3e6c174084e4ce2056f3d60 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/36c460ab5de16c3fd8d0dbf15b52b176 b/me-seed-anneal_25000_8/ocdbt.process_0/d/36c460ab5de16c3fd8d0dbf15b52b176
new file mode 100644
index 0000000000000000000000000000000000000000..10818a35896b49e754401e76e9d118699c216539
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/36c460ab5de16c3fd8d0dbf15b52b176 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/37fd220cf725e95a6c22939dd21604c9 b/me-seed-anneal_25000_8/ocdbt.process_0/d/37fd220cf725e95a6c22939dd21604c9
new file mode 100644
index 0000000000000000000000000000000000000000..dc2d9528cb95ede9872e984b25fcd9b29d11b2cf
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/37fd220cf725e95a6c22939dd21604c9 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/395fbd310a7bbcfaa9e7df0bb230ea5f b/me-seed-anneal_25000_8/ocdbt.process_0/d/395fbd310a7bbcfaa9e7df0bb230ea5f
new file mode 100644
index 0000000000000000000000000000000000000000..554b032263263858d8f1d046ee61a87678bf4dd7
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/395fbd310a7bbcfaa9e7df0bb230ea5f differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/3ea316322c2c7ccdc3a4bd7dec60943c b/me-seed-anneal_25000_8/ocdbt.process_0/d/3ea316322c2c7ccdc3a4bd7dec60943c
new file mode 100644
index 0000000000000000000000000000000000000000..803a6bdd35f3668ccce8db846b57ae7c9793fd7f
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/3ea316322c2c7ccdc3a4bd7dec60943c differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/449a5462932109d3f4057a73a209f16b b/me-seed-anneal_25000_8/ocdbt.process_0/d/449a5462932109d3f4057a73a209f16b
new file mode 100644
index 0000000000000000000000000000000000000000..ce5a69c0a7361bbfa638b25bb88cfe5e5e03336a
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/449a5462932109d3f4057a73a209f16b differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/4f6c6c79f135736cc26d033cbb28cbaf b/me-seed-anneal_25000_8/ocdbt.process_0/d/4f6c6c79f135736cc26d033cbb28cbaf
new file mode 100644
index 0000000000000000000000000000000000000000..7a7c2449d7348d6695d8943838e872bbdfb94063
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/4f6c6c79f135736cc26d033cbb28cbaf differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/54c0a82c9767ddf5a61e42d574d226e5 b/me-seed-anneal_25000_8/ocdbt.process_0/d/54c0a82c9767ddf5a61e42d574d226e5
new file mode 100644
index 0000000000000000000000000000000000000000..0a1593e72c7988236b346f24544aebf6c14c31b2
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/54c0a82c9767ddf5a61e42d574d226e5 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/55eca8ea23c76615fdbe66bd0c707697 b/me-seed-anneal_25000_8/ocdbt.process_0/d/55eca8ea23c76615fdbe66bd0c707697
new file mode 100644
index 0000000000000000000000000000000000000000..31afe63c58af283fcac35250f0b9026b3078ed5d
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/55eca8ea23c76615fdbe66bd0c707697 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/58028a6a01de9682d7fb43f00896e342 b/me-seed-anneal_25000_8/ocdbt.process_0/d/58028a6a01de9682d7fb43f00896e342
new file mode 100644
index 0000000000000000000000000000000000000000..d461070c9615b02b998498c53503ac9c23efc3fe
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/58028a6a01de9682d7fb43f00896e342 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/5e9185782d5c9927f8c54c96e63fcff7 b/me-seed-anneal_25000_8/ocdbt.process_0/d/5e9185782d5c9927f8c54c96e63fcff7
new file mode 100644
index 0000000000000000000000000000000000000000..66c10b347f8d4e7b4e126e291a7f6c453b6118d2
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/5e9185782d5c9927f8c54c96e63fcff7 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/5f1afe917d1713abc8f65693b51e9469 b/me-seed-anneal_25000_8/ocdbt.process_0/d/5f1afe917d1713abc8f65693b51e9469
new file mode 100644
index 0000000000000000000000000000000000000000..9268c3471081572346855312b5494e99276c8c26
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/5f1afe917d1713abc8f65693b51e9469 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/66169ba653fa977241780f1de73151eb b/me-seed-anneal_25000_8/ocdbt.process_0/d/66169ba653fa977241780f1de73151eb
new file mode 100644
index 0000000000000000000000000000000000000000..29b0649ac0c2e8dc5bef04ba142971125b6849dd
--- /dev/null
+++ b/me-seed-anneal_25000_8/ocdbt.process_0/d/66169ba653fa977241780f1de73151eb
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e045577f465d67fb8b05220dd1cfce6ef36ced82a7c48b084ea8a89a6d07b45e
+size 2675060
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/66359981d9799c7d98f3bab5fa953e7e b/me-seed-anneal_25000_8/ocdbt.process_0/d/66359981d9799c7d98f3bab5fa953e7e
new file mode 100644
index 0000000000000000000000000000000000000000..d872462676f28eef35d6cd1650e00a813e0c466d
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/66359981d9799c7d98f3bab5fa953e7e differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/684845605597a43ed025dcded4313ccd b/me-seed-anneal_25000_8/ocdbt.process_0/d/684845605597a43ed025dcded4313ccd
new file mode 100644
index 0000000000000000000000000000000000000000..b360328d4a228cc7d80ccf080392b7615999eed3
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/684845605597a43ed025dcded4313ccd differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/6f9020e172ecd2386ddb884d8691b01a b/me-seed-anneal_25000_8/ocdbt.process_0/d/6f9020e172ecd2386ddb884d8691b01a
new file mode 100644
index 0000000000000000000000000000000000000000..e29b630d272e74cb690468a34da45fafa7469ded
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/6f9020e172ecd2386ddb884d8691b01a differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/754d0eb45c41295fd9e85974b3022f97 b/me-seed-anneal_25000_8/ocdbt.process_0/d/754d0eb45c41295fd9e85974b3022f97
new file mode 100644
index 0000000000000000000000000000000000000000..d0cdaef16b373391c5d87f4df379f8c3c562119c
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/754d0eb45c41295fd9e85974b3022f97 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/7ddda7b1dc419f4f783932e3f57c4f2c b/me-seed-anneal_25000_8/ocdbt.process_0/d/7ddda7b1dc419f4f783932e3f57c4f2c
new file mode 100644
index 0000000000000000000000000000000000000000..53ec735490073f51aa89ed8b7cec445160dc6122
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/7ddda7b1dc419f4f783932e3f57c4f2c differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/7e827cd1d74342d48c4a27453c7321b0 b/me-seed-anneal_25000_8/ocdbt.process_0/d/7e827cd1d74342d48c4a27453c7321b0
new file mode 100644
index 0000000000000000000000000000000000000000..559847b53d4bab9da223d71c912210bc9c65b91e
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/7e827cd1d74342d48c4a27453c7321b0 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/862f8ca374670ff1c7e10b09270769c3 b/me-seed-anneal_25000_8/ocdbt.process_0/d/862f8ca374670ff1c7e10b09270769c3
new file mode 100644
index 0000000000000000000000000000000000000000..d822f019405dd8bb6011bc14b3e61d17faa499c1
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/862f8ca374670ff1c7e10b09270769c3 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/873c9fd29bb44c9d8212785a03d1b866 b/me-seed-anneal_25000_8/ocdbt.process_0/d/873c9fd29bb44c9d8212785a03d1b866
new file mode 100644
index 0000000000000000000000000000000000000000..4567136a80a37eee4919fefdafbd230d909407a8
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/873c9fd29bb44c9d8212785a03d1b866 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/8d2e2d63850079028cab465c3c7c5df4 b/me-seed-anneal_25000_8/ocdbt.process_0/d/8d2e2d63850079028cab465c3c7c5df4
new file mode 100644
index 0000000000000000000000000000000000000000..6c72a68beee8abfd5122f7e794554e5525de3696
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/8d2e2d63850079028cab465c3c7c5df4 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/8f51e8a7e47c1f6dd2bb0e31b416f13d b/me-seed-anneal_25000_8/ocdbt.process_0/d/8f51e8a7e47c1f6dd2bb0e31b416f13d
new file mode 100644
index 0000000000000000000000000000000000000000..958d85b03fe6a6dbeb21728b36cf974013c298ce
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/8f51e8a7e47c1f6dd2bb0e31b416f13d differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/a242e476e90b40485e2d4e1a468efa4e b/me-seed-anneal_25000_8/ocdbt.process_0/d/a242e476e90b40485e2d4e1a468efa4e
new file mode 100644
index 0000000000000000000000000000000000000000..500eb19b7716afee1a0c19b905bb79ae0420657a
--- /dev/null
+++ b/me-seed-anneal_25000_8/ocdbt.process_0/d/a242e476e90b40485e2d4e1a468efa4e
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:28400bb18ff264c39c6510ea1708a15298c01b4c96e1b969e36e82c870e69f41
+size 142878
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/a9710c596fe5521ce1afa79ce7f0d5b8 b/me-seed-anneal_25000_8/ocdbt.process_0/d/a9710c596fe5521ce1afa79ce7f0d5b8
new file mode 100644
index 0000000000000000000000000000000000000000..1870cb5f40f36c3ff0aadb1df311488c751ef6aa
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/a9710c596fe5521ce1afa79ce7f0d5b8 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/ab959637fedfc91e91ab1d0d10f5e13d b/me-seed-anneal_25000_8/ocdbt.process_0/d/ab959637fedfc91e91ab1d0d10f5e13d
new file mode 100644
index 0000000000000000000000000000000000000000..c3b4f358faf5a3bcbcd4034d7672109209bb736e
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/ab959637fedfc91e91ab1d0d10f5e13d differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/b0841a829b9fd46e3d40319ca65c48f3 b/me-seed-anneal_25000_8/ocdbt.process_0/d/b0841a829b9fd46e3d40319ca65c48f3
new file mode 100644
index 0000000000000000000000000000000000000000..3474d655c7868c102b7e97fd9c89f3113501638c
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/b0841a829b9fd46e3d40319ca65c48f3 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/b913624855aeb9a98a3fac97f240bd42 b/me-seed-anneal_25000_8/ocdbt.process_0/d/b913624855aeb9a98a3fac97f240bd42
new file mode 100644
index 0000000000000000000000000000000000000000..28d6d5ee3e82f6420368027f5842ef632f29309a
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/b913624855aeb9a98a3fac97f240bd42 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/bc77d364f449d624aa7f359a91762310 b/me-seed-anneal_25000_8/ocdbt.process_0/d/bc77d364f449d624aa7f359a91762310
new file mode 100644
index 0000000000000000000000000000000000000000..769070f6369f8851e517e59a30edb5f069d9dafe
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/bc77d364f449d624aa7f359a91762310 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/bcef89c346c7f343565ad259c272e9bd b/me-seed-anneal_25000_8/ocdbt.process_0/d/bcef89c346c7f343565ad259c272e9bd
new file mode 100644
index 0000000000000000000000000000000000000000..3cacf801f63d7dd73431a9e8aca3744914fda662
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/bcef89c346c7f343565ad259c272e9bd differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/c528059cd1d8a45d1b25a836ef7f8ec2 b/me-seed-anneal_25000_8/ocdbt.process_0/d/c528059cd1d8a45d1b25a836ef7f8ec2
new file mode 100644
index 0000000000000000000000000000000000000000..966ea803df57a24fb7d36b8027ce2b049915ccd1
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/c528059cd1d8a45d1b25a836ef7f8ec2 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/c703312ed03504950c91a36f9c953f08 b/me-seed-anneal_25000_8/ocdbt.process_0/d/c703312ed03504950c91a36f9c953f08
new file mode 100644
index 0000000000000000000000000000000000000000..8bfb332eab7d9bbddbf9c607d31def024d282dae
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/c703312ed03504950c91a36f9c953f08 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/d5a7ba333c865e603bb5668d48eb8066 b/me-seed-anneal_25000_8/ocdbt.process_0/d/d5a7ba333c865e603bb5668d48eb8066
new file mode 100644
index 0000000000000000000000000000000000000000..2fe8f207f28432fe17c7b34a2bc03c4bd1a4df6f
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/d5a7ba333c865e603bb5668d48eb8066 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/d93c92aab33793bdc7b4c207ad1b12ab b/me-seed-anneal_25000_8/ocdbt.process_0/d/d93c92aab33793bdc7b4c207ad1b12ab
new file mode 100644
index 0000000000000000000000000000000000000000..ee67f22c8a5b88325a4d5c276382c86d108cd1f0
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/d93c92aab33793bdc7b4c207ad1b12ab differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/da570605a800e65d8ad7c36f5183b4fc b/me-seed-anneal_25000_8/ocdbt.process_0/d/da570605a800e65d8ad7c36f5183b4fc
new file mode 100644
index 0000000000000000000000000000000000000000..45cf8deb7f584ff0eae31c442461b90de87a9e38
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/da570605a800e65d8ad7c36f5183b4fc differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/def363cadaf024529b5761f1ecff7b68 b/me-seed-anneal_25000_8/ocdbt.process_0/d/def363cadaf024529b5761f1ecff7b68
new file mode 100644
index 0000000000000000000000000000000000000000..751388b9cc36173d92adedd4718aa785038bedf0
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/def363cadaf024529b5761f1ecff7b68 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/df18c495f53d51e8fb34e59382faec78 b/me-seed-anneal_25000_8/ocdbt.process_0/d/df18c495f53d51e8fb34e59382faec78
new file mode 100644
index 0000000000000000000000000000000000000000..6d000c7ca566a3d09741dbdb264b130ef18c4f53
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/df18c495f53d51e8fb34e59382faec78 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/e5caaa76ed2e1af0c0e651f9899c4e75 b/me-seed-anneal_25000_8/ocdbt.process_0/d/e5caaa76ed2e1af0c0e651f9899c4e75
new file mode 100644
index 0000000000000000000000000000000000000000..99e9e5db40f2e38bb57d0a22b3eeba630b1837c9
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/e5caaa76ed2e1af0c0e651f9899c4e75 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/ed2ae4eadbd5abaea8e9cc59055eed98 b/me-seed-anneal_25000_8/ocdbt.process_0/d/ed2ae4eadbd5abaea8e9cc59055eed98
new file mode 100644
index 0000000000000000000000000000000000000000..9ca6eb61aeacc228a48cfa60600e0a55509e14f2
--- /dev/null
+++ b/me-seed-anneal_25000_8/ocdbt.process_0/d/ed2ae4eadbd5abaea8e9cc59055eed98
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f4d2ff10ba6930e26450997cf17bc155d60356f7cf1183c6109b98d367c833d4
+size 9056256
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/eee6bd35a348721351eb051502b81ace b/me-seed-anneal_25000_8/ocdbt.process_0/d/eee6bd35a348721351eb051502b81ace
new file mode 100644
index 0000000000000000000000000000000000000000..533e793152853b7fe2d05314410a35122adb2efc
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/eee6bd35a348721351eb051502b81ace differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/efd85ba6867fe7517f7fefc52ce8eac0 b/me-seed-anneal_25000_8/ocdbt.process_0/d/efd85ba6867fe7517f7fefc52ce8eac0
new file mode 100644
index 0000000000000000000000000000000000000000..04df9d53589dfe32ead0d0626ad9367f4785af22
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/efd85ba6867fe7517f7fefc52ce8eac0 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/f08ceaa1d45887e48b31afe502c44ea1 b/me-seed-anneal_25000_8/ocdbt.process_0/d/f08ceaa1d45887e48b31afe502c44ea1
new file mode 100644
index 0000000000000000000000000000000000000000..427821679c4df4494fd543cab50d4562c91adc7b
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/f08ceaa1d45887e48b31afe502c44ea1 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/f0a7ca7911c6023e7a89b9f11d3abc1b b/me-seed-anneal_25000_8/ocdbt.process_0/d/f0a7ca7911c6023e7a89b9f11d3abc1b
new file mode 100644
index 0000000000000000000000000000000000000000..7511c7ad6e03d63c21c5f761785d81053c9e5474
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/f0a7ca7911c6023e7a89b9f11d3abc1b differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/f19baac0b3158a145c6d0a3f92591a97 b/me-seed-anneal_25000_8/ocdbt.process_0/d/f19baac0b3158a145c6d0a3f92591a97
new file mode 100644
index 0000000000000000000000000000000000000000..38fe9474f61968141d143549a6d7d5a18b9d7c4e
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/f19baac0b3158a145c6d0a3f92591a97 differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/f30fd90dbe43cf18050c09cc61d7891f b/me-seed-anneal_25000_8/ocdbt.process_0/d/f30fd90dbe43cf18050c09cc61d7891f
new file mode 100644
index 0000000000000000000000000000000000000000..ae25c1b1733587cb47a80f8686cb1221b7cb1cc8
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/f30fd90dbe43cf18050c09cc61d7891f differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/f3ee743bbd904a50313ab0bb0b44435e b/me-seed-anneal_25000_8/ocdbt.process_0/d/f3ee743bbd904a50313ab0bb0b44435e
new file mode 100644
index 0000000000000000000000000000000000000000..4d0eff4eb391186d9077ac85dcd0660e58a7bf43
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/f3ee743bbd904a50313ab0bb0b44435e differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/f3eff486d97fb3ed969771bea8c08f9e b/me-seed-anneal_25000_8/ocdbt.process_0/d/f3eff486d97fb3ed969771bea8c08f9e
new file mode 100644
index 0000000000000000000000000000000000000000..d8a75412525d939524371040607a9945c7591366
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/f3eff486d97fb3ed969771bea8c08f9e differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/f4df26ebeeda30cb1cde15a19c867031 b/me-seed-anneal_25000_8/ocdbt.process_0/d/f4df26ebeeda30cb1cde15a19c867031
new file mode 100644
index 0000000000000000000000000000000000000000..830e74cd3fdbc7eda85a56886d6aa8c9268fb6b5
--- /dev/null
+++ b/me-seed-anneal_25000_8/ocdbt.process_0/d/f4df26ebeeda30cb1cde15a19c867031
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e2e3fae6b599d35abefeafe86f90c0988d1966233f0f2d6d1d4ae3e523f5ce8a
+size 47693824
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/d/f92f345bb17bd198a9309c4442c6980b b/me-seed-anneal_25000_8/ocdbt.process_0/d/f92f345bb17bd198a9309c4442c6980b
new file mode 100644
index 0000000000000000000000000000000000000000..e542ab1f12400f71dc952e8a2ca83f3f60b5b73f
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/d/f92f345bb17bd198a9309c4442c6980b differ
diff --git a/me-seed-anneal_25000_8/ocdbt.process_0/manifest.ocdbt b/me-seed-anneal_25000_8/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..33503105977bfbe012235e7de0ae7a279a4f23f2
Binary files /dev/null and b/me-seed-anneal_25000_8/ocdbt.process_0/manifest.ocdbt differ