diff --git a/.gitattributes b/.gitattributes
index 41cec71ce7e7c61ba4bb79dc149169e60d99b3de..8312560cde75085559a80d48bb1c14b42708e7ad 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -82,3 +82,8 @@ me-seed-anneal_25000_8/ocdbt.process_0/d/66169ba653fa977241780f1de73151eb filter
 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
+me-seed_25000_2/d/bfe8856423c2a9ebd541a5d25a5b3fe8 filter=lfs diff=lfs merge=lfs -text
+me-seed_25000_2/ocdbt.process_0/d/620d379dbbb8e4b1aa40fac633ead069 filter=lfs diff=lfs merge=lfs -text
+me-seed_25000_2/ocdbt.process_0/d/973b07218be650bf6df4db1b19f35382 filter=lfs diff=lfs merge=lfs -text
+me-seed_25000_2/ocdbt.process_0/d/caca59e162b50bcc68be044ae66aa676 filter=lfs diff=lfs merge=lfs -text
+me-seed_25000_2/ocdbt.process_0/d/fdbae4faad18c96b7dce3ccf61d52b06 filter=lfs diff=lfs merge=lfs -text
diff --git a/me-seed_25000_2/_CHECKPOINT_METADATA b/me-seed_25000_2/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..66b3481413ef5257025b39ee1546b37fa0031a01
--- /dev/null
+++ b/me-seed_25000_2/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786427690186667025, "commit_timestamp_nsecs": 1786427699847862198, "custom_metadata": {}}
\ No newline at end of file
diff --git a/me-seed_25000_2/_METADATA b/me-seed_25000_2/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..34e12fdaa3372cfbfc072b292f31e35a9fd1d8b5
--- /dev/null
+++ b/me-seed_25000_2/_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_25000_2/_sharding b/me-seed_25000_2/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..a41c54a72b12bc2f8c6f57135c83c260b8a7b589
--- /dev/null
+++ b/me-seed_25000_2/_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_25000_2/array_metadatas/process_0 b/me-seed_25000_2/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..cab9b2c86ee1e3c5eddae4d89c33ae2629719b66
--- /dev/null
+++ b/me-seed_25000_2/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_25000_2/d/bfe8856423c2a9ebd541a5d25a5b3fe8 b/me-seed_25000_2/d/bfe8856423c2a9ebd541a5d25a5b3fe8
new file mode 100644
index 0000000000000000000000000000000000000000..e79c44b9121ba44d8953584cfa733d2040d40529
--- /dev/null
+++ b/me-seed_25000_2/d/bfe8856423c2a9ebd541a5d25a5b3fe8
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:460a1f965211c8306b06f6c3deeceb52f8c14681e29bb92d5c31ffac11f970e1
+size 842090
diff --git a/me-seed_25000_2/manifest.ocdbt b/me-seed_25000_2/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..5f23c3777f91b1605aab94cf3ac4d102039dbf5f
Binary files /dev/null and b/me-seed_25000_2/manifest.ocdbt differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/00a1afee8f9ac7e065d0da5afca9c24c b/me-seed_25000_2/ocdbt.process_0/d/00a1afee8f9ac7e065d0da5afca9c24c
new file mode 100644
index 0000000000000000000000000000000000000000..42217ea207080fc8e3fdf6c9e58bfadfdf229cb2
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/00a1afee8f9ac7e065d0da5afca9c24c differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/02b21e3b0dde0afa911a9f21f4793698 b/me-seed_25000_2/ocdbt.process_0/d/02b21e3b0dde0afa911a9f21f4793698
new file mode 100644
index 0000000000000000000000000000000000000000..b89fd0cda012e384079c00609a9301e2f2a83587
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/02b21e3b0dde0afa911a9f21f4793698 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/032f94d60b3fc742cd3f80022cd025eb b/me-seed_25000_2/ocdbt.process_0/d/032f94d60b3fc742cd3f80022cd025eb
new file mode 100644
index 0000000000000000000000000000000000000000..b216fdd228b3e1fba71bba02215db84d7bfc384b
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/032f94d60b3fc742cd3f80022cd025eb differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/065c8e5408eca4c0cbe258b2f64b2cc2 b/me-seed_25000_2/ocdbt.process_0/d/065c8e5408eca4c0cbe258b2f64b2cc2
new file mode 100644
index 0000000000000000000000000000000000000000..cb2d20f042331077cbcfc0ada0287d1b1c82500e
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/065c8e5408eca4c0cbe258b2f64b2cc2 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/07c5bdb7f5c39ddba031d671651ec1b9 b/me-seed_25000_2/ocdbt.process_0/d/07c5bdb7f5c39ddba031d671651ec1b9
new file mode 100644
index 0000000000000000000000000000000000000000..cbe9db594a1ce1f518bbd2971e664a19294f8177
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/07c5bdb7f5c39ddba031d671651ec1b9 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/0c333dc946dfa0672fabba66c4ee9f08 b/me-seed_25000_2/ocdbt.process_0/d/0c333dc946dfa0672fabba66c4ee9f08
new file mode 100644
index 0000000000000000000000000000000000000000..76155492f0f5bfba80949f5e74052531d09e1966
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/0c333dc946dfa0672fabba66c4ee9f08 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/0d969ff01bf97f466d25f2c4a78a4de6 b/me-seed_25000_2/ocdbt.process_0/d/0d969ff01bf97f466d25f2c4a78a4de6
new file mode 100644
index 0000000000000000000000000000000000000000..0f1a037013a9bd348e5ea32d005e6d144579f5df
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/0d969ff01bf97f466d25f2c4a78a4de6 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/125513475e84c3b4e68ea191c800527d b/me-seed_25000_2/ocdbt.process_0/d/125513475e84c3b4e68ea191c800527d
new file mode 100644
index 0000000000000000000000000000000000000000..51b5aaaf29bac6ab6844b79f593446d753f65314
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/125513475e84c3b4e68ea191c800527d differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/1db6a45f3f60cc01d158598e72d4fa58 b/me-seed_25000_2/ocdbt.process_0/d/1db6a45f3f60cc01d158598e72d4fa58
new file mode 100644
index 0000000000000000000000000000000000000000..48e8c4c2d5a6f518118206a3da62c0a0a55ce903
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/1db6a45f3f60cc01d158598e72d4fa58 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/1e53865483632d69a0263d795ed59304 b/me-seed_25000_2/ocdbt.process_0/d/1e53865483632d69a0263d795ed59304
new file mode 100644
index 0000000000000000000000000000000000000000..0c4860afc7ee0341dd75cc6be55c935505fe8a9c
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/1e53865483632d69a0263d795ed59304 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/20229f6c10d0bdd623d4bc1a69db5444 b/me-seed_25000_2/ocdbt.process_0/d/20229f6c10d0bdd623d4bc1a69db5444
new file mode 100644
index 0000000000000000000000000000000000000000..6d983a07c0210e4e62f876184bbe23a3d3a1fe06
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/20229f6c10d0bdd623d4bc1a69db5444 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/20904ef5d74f9e0e8c9876f59af248c6 b/me-seed_25000_2/ocdbt.process_0/d/20904ef5d74f9e0e8c9876f59af248c6
new file mode 100644
index 0000000000000000000000000000000000000000..a3dd6844403e995ac09f8790930a306fda6370fd
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/20904ef5d74f9e0e8c9876f59af248c6 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/20b76276971809c2c20c807d996a9796 b/me-seed_25000_2/ocdbt.process_0/d/20b76276971809c2c20c807d996a9796
new file mode 100644
index 0000000000000000000000000000000000000000..2e391373336d27cf1716cbdbaf2b138d3028fcde
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/20b76276971809c2c20c807d996a9796 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/21e4534352bc024573edf1483b2ec02f b/me-seed_25000_2/ocdbt.process_0/d/21e4534352bc024573edf1483b2ec02f
new file mode 100644
index 0000000000000000000000000000000000000000..29cfa624254e2a14f2e27910981ce61e277777aa
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/21e4534352bc024573edf1483b2ec02f differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/27769a252b8f1d9d37ecd9081d4ca6d4 b/me-seed_25000_2/ocdbt.process_0/d/27769a252b8f1d9d37ecd9081d4ca6d4
new file mode 100644
index 0000000000000000000000000000000000000000..3078e3c2ba1c09df0e987571d7d4225acb92e1fd
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/27769a252b8f1d9d37ecd9081d4ca6d4 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/2ed879221c889c4994efdb6057efcf18 b/me-seed_25000_2/ocdbt.process_0/d/2ed879221c889c4994efdb6057efcf18
new file mode 100644
index 0000000000000000000000000000000000000000..89b5abe6a79e48d3e017c65d1437706d3b91cfe3
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/2ed879221c889c4994efdb6057efcf18 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/30631de89b43c734f3cac384ec6c5d00 b/me-seed_25000_2/ocdbt.process_0/d/30631de89b43c734f3cac384ec6c5d00
new file mode 100644
index 0000000000000000000000000000000000000000..5db1e8c991ce86240131a45906e2b058d3f58ddc
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/30631de89b43c734f3cac384ec6c5d00 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/3ed8c09f742c5566beac4b3e097fc517 b/me-seed_25000_2/ocdbt.process_0/d/3ed8c09f742c5566beac4b3e097fc517
new file mode 100644
index 0000000000000000000000000000000000000000..9d7e95de43504070f6d49420cda5e9178db77d56
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/3ed8c09f742c5566beac4b3e097fc517 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/413071d0836b85b93a94d8647bfaec09 b/me-seed_25000_2/ocdbt.process_0/d/413071d0836b85b93a94d8647bfaec09
new file mode 100644
index 0000000000000000000000000000000000000000..47fddd0f8149929a85c0534ee58bd3baf22e7595
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/413071d0836b85b93a94d8647bfaec09 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/435411226026ecd81a7b7209142b06c6 b/me-seed_25000_2/ocdbt.process_0/d/435411226026ecd81a7b7209142b06c6
new file mode 100644
index 0000000000000000000000000000000000000000..307d3c3ae568af8213a29385e3e51d2fd2c188c1
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/435411226026ecd81a7b7209142b06c6 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/43f8eeccd80d01bcb39e7004dd7a7844 b/me-seed_25000_2/ocdbt.process_0/d/43f8eeccd80d01bcb39e7004dd7a7844
new file mode 100644
index 0000000000000000000000000000000000000000..5e714e946b785d600542ad3c94041783315ea69e
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/43f8eeccd80d01bcb39e7004dd7a7844 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/48d4389248b06af226bc8f66787a437b b/me-seed_25000_2/ocdbt.process_0/d/48d4389248b06af226bc8f66787a437b
new file mode 100644
index 0000000000000000000000000000000000000000..47fddcf53fe4178c4f9533f4fded7d3c907d8986
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/48d4389248b06af226bc8f66787a437b differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/5d32d3ee758377aff304829805c15339 b/me-seed_25000_2/ocdbt.process_0/d/5d32d3ee758377aff304829805c15339
new file mode 100644
index 0000000000000000000000000000000000000000..bc25e9e152cb8e7dc21d448f7e7cb06e0677d61b
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/5d32d3ee758377aff304829805c15339 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/5efa127eb03478f677adac8e06a5dd4e b/me-seed_25000_2/ocdbt.process_0/d/5efa127eb03478f677adac8e06a5dd4e
new file mode 100644
index 0000000000000000000000000000000000000000..8600039369eb0fd0c680e9b33499945475e76589
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/5efa127eb03478f677adac8e06a5dd4e differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/620d379dbbb8e4b1aa40fac633ead069 b/me-seed_25000_2/ocdbt.process_0/d/620d379dbbb8e4b1aa40fac633ead069
new file mode 100644
index 0000000000000000000000000000000000000000..b40e34ed9703d3292bde36499e8cc421e81e0be0
--- /dev/null
+++ b/me-seed_25000_2/ocdbt.process_0/d/620d379dbbb8e4b1aa40fac633ead069
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:061f9a155a0426d43911ef32a45244717edb32d31766c8419465c64ad001d819
+size 20512768
diff --git a/me-seed_25000_2/ocdbt.process_0/d/63da8d0727c4cfa86cf8b34d3a6ce943 b/me-seed_25000_2/ocdbt.process_0/d/63da8d0727c4cfa86cf8b34d3a6ce943
new file mode 100644
index 0000000000000000000000000000000000000000..bc505488b13297f185068f45ee3a993860f668cc
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/63da8d0727c4cfa86cf8b34d3a6ce943 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/6541a2c69d679950f10c84eeeb7fbbeb b/me-seed_25000_2/ocdbt.process_0/d/6541a2c69d679950f10c84eeeb7fbbeb
new file mode 100644
index 0000000000000000000000000000000000000000..0849fb27f393ff6b0622a683cfa50bf60e4799bc
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/6541a2c69d679950f10c84eeeb7fbbeb differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/6cf8f6c89e0d229ada0922d7f37ff470 b/me-seed_25000_2/ocdbt.process_0/d/6cf8f6c89e0d229ada0922d7f37ff470
new file mode 100644
index 0000000000000000000000000000000000000000..80ceaec5c0aa7fa767408a668353bf505096edb8
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/6cf8f6c89e0d229ada0922d7f37ff470 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/6d0e27253ad7c24f30746df8e198c83d b/me-seed_25000_2/ocdbt.process_0/d/6d0e27253ad7c24f30746df8e198c83d
new file mode 100644
index 0000000000000000000000000000000000000000..f9b7be835c4f4e9b384dfb5f177130303762f984
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/6d0e27253ad7c24f30746df8e198c83d differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/76b5b529a71ea18520f50582c3a9a6a0 b/me-seed_25000_2/ocdbt.process_0/d/76b5b529a71ea18520f50582c3a9a6a0
new file mode 100644
index 0000000000000000000000000000000000000000..8d7851ebd26b6668bda54efe2af49fe53b892abf
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/76b5b529a71ea18520f50582c3a9a6a0 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/7b8255438734fa71b5ae82502130e225 b/me-seed_25000_2/ocdbt.process_0/d/7b8255438734fa71b5ae82502130e225
new file mode 100644
index 0000000000000000000000000000000000000000..1c4590d2971eaf09b6c8e9e98b434f261ca5ae18
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/7b8255438734fa71b5ae82502130e225 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/7dd5cb033ab603b14eec6084544882b5 b/me-seed_25000_2/ocdbt.process_0/d/7dd5cb033ab603b14eec6084544882b5
new file mode 100644
index 0000000000000000000000000000000000000000..113cd6ffc05295bf3b5f366eace8038684dcd2ea
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/7dd5cb033ab603b14eec6084544882b5 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/973b07218be650bf6df4db1b19f35382 b/me-seed_25000_2/ocdbt.process_0/d/973b07218be650bf6df4db1b19f35382
new file mode 100644
index 0000000000000000000000000000000000000000..56834addc647b1765898089c8954319eccfa1550
--- /dev/null
+++ b/me-seed_25000_2/ocdbt.process_0/d/973b07218be650bf6df4db1b19f35382
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2a994ee246e90e5d1c00eb08e1c16d6e6cf834366943ff1b1f863494696eb9b4
+size 15618048
diff --git a/me-seed_25000_2/ocdbt.process_0/d/9f1ac7060180cadfde7f0acc86e3b487 b/me-seed_25000_2/ocdbt.process_0/d/9f1ac7060180cadfde7f0acc86e3b487
new file mode 100644
index 0000000000000000000000000000000000000000..ee66fb0b47c41930b559bd728022017db40cfa4e
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/9f1ac7060180cadfde7f0acc86e3b487 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/a2cf9fd344939d771b1ed2c0dcf5f53c b/me-seed_25000_2/ocdbt.process_0/d/a2cf9fd344939d771b1ed2c0dcf5f53c
new file mode 100644
index 0000000000000000000000000000000000000000..1831535e00c1759946a17fc4a408b8e932512949
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/a2cf9fd344939d771b1ed2c0dcf5f53c differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/ab63c9de9866fba1d28dd1db89f9685d b/me-seed_25000_2/ocdbt.process_0/d/ab63c9de9866fba1d28dd1db89f9685d
new file mode 100644
index 0000000000000000000000000000000000000000..cc80ad8df5b8b1b26e9a52d15ad12187daea8da6
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/ab63c9de9866fba1d28dd1db89f9685d differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/b752d8d1f51f7990780577712bbf2b0f b/me-seed_25000_2/ocdbt.process_0/d/b752d8d1f51f7990780577712bbf2b0f
new file mode 100644
index 0000000000000000000000000000000000000000..07a0a0da4077d7da850f31a1b3f84c79fc2dac8b
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/b752d8d1f51f7990780577712bbf2b0f differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/bcad885834d3dd9f54dedf5933b234b6 b/me-seed_25000_2/ocdbt.process_0/d/bcad885834d3dd9f54dedf5933b234b6
new file mode 100644
index 0000000000000000000000000000000000000000..ce5a69c0a7361bbfa638b25bb88cfe5e5e03336a
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/bcad885834d3dd9f54dedf5933b234b6 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/c60c95bcf2bf58ade412624ce41bc60f b/me-seed_25000_2/ocdbt.process_0/d/c60c95bcf2bf58ade412624ce41bc60f
new file mode 100644
index 0000000000000000000000000000000000000000..81b8c1d714467ce140c282933bbc3830523a56a0
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/c60c95bcf2bf58ade412624ce41bc60f differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/caca59e162b50bcc68be044ae66aa676 b/me-seed_25000_2/ocdbt.process_0/d/caca59e162b50bcc68be044ae66aa676
new file mode 100644
index 0000000000000000000000000000000000000000..ccfb3ede4498d456387cc22dbc83965512901c00
--- /dev/null
+++ b/me-seed_25000_2/ocdbt.process_0/d/caca59e162b50bcc68be044ae66aa676
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:75592d2b94a2ac3bfe04ec8dbcf99809b4201f4b00ae932f2b1d39a7aacc5910
+size 233366
diff --git a/me-seed_25000_2/ocdbt.process_0/d/cdbd6ee15386c5fa55bec3594ac81bd1 b/me-seed_25000_2/ocdbt.process_0/d/cdbd6ee15386c5fa55bec3594ac81bd1
new file mode 100644
index 0000000000000000000000000000000000000000..03edfe9137a91eca69e6eede15d2eab9e958c137
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/cdbd6ee15386c5fa55bec3594ac81bd1 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/d60e909aec93a6ac75390f9a774a3e93 b/me-seed_25000_2/ocdbt.process_0/d/d60e909aec93a6ac75390f9a774a3e93
new file mode 100644
index 0000000000000000000000000000000000000000..408293c9425f55d0e518cd0f73e8ede5ceadf5fe
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/d60e909aec93a6ac75390f9a774a3e93 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/d9800a743589dc76f743f919b24473a4 b/me-seed_25000_2/ocdbt.process_0/d/d9800a743589dc76f743f919b24473a4
new file mode 100644
index 0000000000000000000000000000000000000000..4a055fb271fd04914ca2bdecd1f4214d7209cf60
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/d9800a743589dc76f743f919b24473a4 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/dace27990e520da8f702f0cf90c04e10 b/me-seed_25000_2/ocdbt.process_0/d/dace27990e520da8f702f0cf90c04e10
new file mode 100644
index 0000000000000000000000000000000000000000..41b0e3ec7dc8c1974e679bee455da41c3356ca34
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/dace27990e520da8f702f0cf90c04e10 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/dba98612adedb8d4b71af624fa714fb2 b/me-seed_25000_2/ocdbt.process_0/d/dba98612adedb8d4b71af624fa714fb2
new file mode 100644
index 0000000000000000000000000000000000000000..951e47a38b4f1e3765d33f9af78e0a999edb6c77
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/dba98612adedb8d4b71af624fa714fb2 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/e0b7b28c099895eec59bbb9f629828c2 b/me-seed_25000_2/ocdbt.process_0/d/e0b7b28c099895eec59bbb9f629828c2
new file mode 100644
index 0000000000000000000000000000000000000000..7a7c2449d7348d6695d8943838e872bbdfb94063
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/e0b7b28c099895eec59bbb9f629828c2 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/e2b991e8b336e0aeb058006eaae043af b/me-seed_25000_2/ocdbt.process_0/d/e2b991e8b336e0aeb058006eaae043af
new file mode 100644
index 0000000000000000000000000000000000000000..04d93de57c4a3010c92820aa1bcbde4ddab4bca0
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/e2b991e8b336e0aeb058006eaae043af differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/fa9757932275109a2c80947d90b76f42 b/me-seed_25000_2/ocdbt.process_0/d/fa9757932275109a2c80947d90b76f42
new file mode 100644
index 0000000000000000000000000000000000000000..098c23ccb245c235ea613df43fb4382b9e11d5cd
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/fa9757932275109a2c80947d90b76f42 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/fd8ad71eacc6674020eacede7110879c b/me-seed_25000_2/ocdbt.process_0/d/fd8ad71eacc6674020eacede7110879c
new file mode 100644
index 0000000000000000000000000000000000000000..21e8cd48f6464b3a546e2de9f962e6c0bdd0a1f5
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/fd8ad71eacc6674020eacede7110879c differ
diff --git a/me-seed_25000_2/ocdbt.process_0/d/fdbae4faad18c96b7dce3ccf61d52b06 b/me-seed_25000_2/ocdbt.process_0/d/fdbae4faad18c96b7dce3ccf61d52b06
new file mode 100644
index 0000000000000000000000000000000000000000..6098cc7791ce740563d12ea4cd2b377272db7718
--- /dev/null
+++ b/me-seed_25000_2/ocdbt.process_0/d/fdbae4faad18c96b7dce3ccf61d52b06
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3cd87d9d5df052d5620437658036f3ad528de56fa4a684644dc97b7fef9a5b8f
+size 171233280
diff --git a/me-seed_25000_2/ocdbt.process_0/d/fdc55c0681e5c29aedefba8e14e0eab5 b/me-seed_25000_2/ocdbt.process_0/d/fdc55c0681e5c29aedefba8e14e0eab5
new file mode 100644
index 0000000000000000000000000000000000000000..20128cde27107ab42eca4a33aedf334d53f87afd
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/d/fdc55c0681e5c29aedefba8e14e0eab5 differ
diff --git a/me-seed_25000_2/ocdbt.process_0/manifest.ocdbt b/me-seed_25000_2/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..5e67bac38d5c612358ef7f0510d6107836393dd9
Binary files /dev/null and b/me-seed_25000_2/ocdbt.process_0/manifest.ocdbt differ