diff --git a/.gitattributes b/.gitattributes
index 2f119661e6fd54546a88740e5e9579dbf768311e..e1f9bdfd3b46e4bb56e4f5b8b4493a54ea67b7cb 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -52,3 +52,7 @@ me-full_25000_5/ocdbt.process_0/d/00af2b5dd8c97f8aa04c90314584796b filter=lfs di
 me-full_25000_5/ocdbt.process_0/d/231e228c0807bec3e5e5dfa0fbe0deb1 filter=lfs diff=lfs merge=lfs -text
 me-full_25000_5/ocdbt.process_0/d/7f413e304642522d0dc6738b225998d7 filter=lfs diff=lfs merge=lfs -text
 me-full_25000_5/ocdbt.process_0/d/a150499bc1dda58bef676984f1c5c136 filter=lfs diff=lfs merge=lfs -text
+me-full_25000_30/d/ed0c67802ee89a318ab7d7957e50af2b filter=lfs diff=lfs merge=lfs -text
+me-full_25000_30/ocdbt.process_0/d/39a6249b85ab5b6b6a6cd7ca1b476d0b filter=lfs diff=lfs merge=lfs -text
+me-full_25000_30/ocdbt.process_0/d/438064f8912286e5c65301d98c3cde35 filter=lfs diff=lfs merge=lfs -text
+me-full_25000_30/ocdbt.process_0/d/d3218fa39ee5f86a0c1cf8bfdc8db625 filter=lfs diff=lfs merge=lfs -text
diff --git a/me-full_25000_30/_CHECKPOINT_METADATA b/me-full_25000_30/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..b109cfd9f41d8b1a926ac79c1f8be38e5327a851
--- /dev/null
+++ b/me-full_25000_30/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1783758754617880452, "commit_timestamp_nsecs": 1783758760632807719, "custom_metadata": {}}
\ No newline at end of file
diff --git a/me-full_25000_30/_METADATA b/me-full_25000_30/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..15bc840afb188b3f12a0cae4c24e9b252b0c1564
--- /dev/null
+++ b/me-full_25000_30/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('0', 'decoder', 'attn_layers', '0', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '0', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('0', 'decoder', 'attn_layers', '0', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '0', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'attn_layers', '1', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '1', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('0', 'decoder', 'attn_layers', '1', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '1', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'dof_gl', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'decoder', 'dof_gl', 'd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'decoder', 'dof_gl', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'decoder', 'dof_gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('0', 'decoder', 'dof_gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'decoder', 'dof_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'dof_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'dof_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('0', 'decoder', 'dof_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'decoder', 'ln_dist_dof', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_dof", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_dist_dof', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_dof", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_dist_pose', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_pose", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_dist_pose', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_pose", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_tmp_dof', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_dof", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_tmp_dof', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_dof", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_tmp_pose', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_pose", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_tmp_pose', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_pose", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'decoder', 'pos_head', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 7, 9]}}, "('0', 'decoder', 'pos_head', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 7, 9]}}, "('0', 'decoder', 'pos_head', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'pos_head', 'grid_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'grid_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'pos_head', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'decoder', 'pos_head', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 64]}}, "('0', 'decoder', 'pos_head', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('0', 'decoder', 'pos_head', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'decoder', 'pos_head', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 1, 64]}}, "('0', 'decoder', 'pose_gl', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'decoder', 'pose_gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pose_gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pose_gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pose_gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 384]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 4]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 384]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'decoder', 'rot_head', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'decoder', 'rot_head', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 7, 9]}}, "('0', 'decoder', 'rot_head', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 7, 9]}}, "('0', 'decoder', 'rot_head', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'rot_head', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'rot_head', 'grid_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'rot_head', 'grid_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'rot_head', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'rot_head', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'decoder', 'rot_head', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 64]}}, "('0', 'decoder', 'rot_head', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('0', 'decoder', 'rot_head', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'decoder', 'rot_head', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 1, 64]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'kinematic_encoding', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('0', 'kinematic_encoding', 'gl', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('0', 'kinematic_encoding', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'kinematic_encoding', 'gl', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'kinematic_encoding', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('0', 'kinematic_encoding', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'base_to_contact')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "base_to_contact", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'dof_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'joint_ranges')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "joint_ranges", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'joint_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "joint_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 6]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'kinematics_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kinematics_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'pose_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'base_to_contact')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "base_to_contact", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'dof_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'joint_ranges')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "joint_ranges", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'joint_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "joint_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 6]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'kinematics_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kinematics_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'pose_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'align_to_approach', '0')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "align_to_approach", "key_type": 2}, {"key": "0", "key_type": 1}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'align_to_approach', '1')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "align_to_approach", "key_type": 2}, {"key": "1", "key_type": 1}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'base_to_contact')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "base_to_contact", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'dof_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [13, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'fingertip_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fingertip_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'fingertip_normals')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fingertip_normals", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'init_pregrasp_joint')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "init_pregrasp_joint", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [12]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'joint_ranges')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "joint_ranges", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [12, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'joint_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "joint_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [12, 6]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'kinematics_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kinematics_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [12, 7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'local_fingertip_contact_positions')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "local_fingertip_contact_positions", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 6, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'pose_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'align_to_approach', '0')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "align_to_approach", "key_type": 2}, {"key": "0", "key_type": 1}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'align_to_approach', '1')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "align_to_approach", "key_type": 2}, {"key": "1", "key_type": 1}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'base_to_contact')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "base_to_contact", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'dof_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [17, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'fingertip_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "fingertip_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'fingertip_normals')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "fingertip_normals", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'init_pregrasp_joint')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "init_pregrasp_joint", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'joint_ranges')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "joint_ranges", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'joint_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "joint_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 6]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'kinematics_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kinematics_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'local_fingertip_contact_positions')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "local_fingertip_contact_positions", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 3, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'pose_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'base_to_contact')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "base_to_contact", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'dof_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [23, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'joint_ranges')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "joint_ranges", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [22, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'joint_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "joint_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [22, 6]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'kinematics_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kinematics_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [22, 7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'pose_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'base_to_contact')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "base_to_contact", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'dof_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'joint_ranges')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "joint_ranges", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'joint_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "joint_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'kinematics_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "kinematics_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'pose_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('0', 'kinematic_encoding', 'ln_dist', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_dist', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_q_dot_sh', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot_sh", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_q_dot_sh', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot_sh", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_time', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_time', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('0', 'kinematic_encoding', 'q_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "q_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'kinematic_encoding', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'kinematic_encoding', 'weight_irrep', 'rtoi')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "weight_irrep", "key_type": 2}, {"key": "rtoi", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 's_proj', 'l_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "s_proj", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [30720]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'layer', 'd_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "d_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1216]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'layer', 'l_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [61440]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_feat', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_feat', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_scalar', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_scalar', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_time', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_time', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '3', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '3', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'q_s_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_s_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'weight_irrep', 'rtoi')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "weight_irrep", "key_type": 2}, {"key": "rtoi", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'biases', 'bias_0')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "biases", "key_type": 2}, {"key": "bias_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '0')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '1')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '2')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "2", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '3')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "3", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '4')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "4", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '5')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "5", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '6')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "6", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '7')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "7", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '8')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "8", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '9')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "9", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '10')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "10", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '11')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "11", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '12')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "12", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '13')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "13", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '14')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "14", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'weight_irrep', 'rtoi')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "weight_irrep", "key_type": 2}, {"key": "rtoi", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 's_proj', 'l_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "s_proj", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [30720]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'layer', 'd_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "d_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1216]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'layer', 'l_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [61440]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_feat', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_feat', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_scalar', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_scalar', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_time', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_time', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '3', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '3', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'q_s_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_s_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'weight_irrep', 'rtoi')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "weight_irrep", "key_type": 2}, {"key": "rtoi", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'biases', 'bias_0')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "biases", "key_type": 2}, {"key": "bias_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '0')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '1')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '2')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "2", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '3')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "3", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '4')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "4", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '5')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "5", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '6')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "6", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '7')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "7", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '8')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "8", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '9')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "9", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '10')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "10", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '11')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "11", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '12')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "12", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '13')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "13", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '14')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "14", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'weight_irrep', 'rtoi')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "weight_irrep", "key_type": 2}, {"key": "rtoi", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 224]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 384]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 192]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 224]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 384]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 192]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 224]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 384]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 192]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'proj', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '0', 'proj', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '0', 'proj', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 32]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 9, 16]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9, 9]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'proj', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'proj', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [9]}}, "('0', 'unet', 'pool_blocks', '1', 'proj', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('0', 'unet', 'weight_irrep', 'rtoi')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_irrep", "key_type": 2}, {"key": "rtoi", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'weight_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'weight_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'weight_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'weight_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'count')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_dist_dof', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_dof", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_dist_dof', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_dof", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_dist_pose', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_pose", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_dist_pose', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_pose", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_tmp_dof', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_dof", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_tmp_dof', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_dof", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_tmp_pose', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_pose", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_tmp_pose', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_pose", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 1, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 384]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 384]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 1, 64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [13, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [17, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [23, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_dist', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_dist', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_q_dot_sh', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot_sh", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_q_dot_sh', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot_sh", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_time', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_time', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'q_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "q_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 's_proj', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "s_proj", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [30720]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'layer', 'd_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "d_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1216]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'layer', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [61440]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_feat', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_feat', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_scalar', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_scalar', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_time', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_time', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '3', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '3', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'q_s_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_s_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'biases', 'bias_0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "biases", "key_type": 2}, {"key": "bias_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '1')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '2')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "2", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '3')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "3", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '4')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "4", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '5')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "5", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '6')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "6", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '7')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "7", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '8')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "8", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '9')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "9", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '10')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "10", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '11')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "11", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '12')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "12", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '13')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "13", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '14')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "14", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 's_proj', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "s_proj", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [30720]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'layer', 'd_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "d_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1216]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'layer', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [61440]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_feat', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_feat', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_scalar', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_scalar', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_time', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_time', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '3', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '3', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'q_s_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_s_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'biases', 'bias_0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "biases", "key_type": 2}, {"key": "bias_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '1')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '2')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "2", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '3')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "3", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '4')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "4", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '5')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "5", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '6')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "6", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '7')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "7", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '8')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "8", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '9')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "9", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '10')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "10", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '11')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "11", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '12')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "12", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '13')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "13", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '14')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "14", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 224]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 224]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 224]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'proj', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'proj', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'proj', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'proj', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'weight_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'weight_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'weight_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'weight_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_dist_dof', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_dof", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_dist_dof', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_dof", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_dist_pose', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_pose", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_dist_pose', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_pose", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_tmp_dof', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_dof", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_tmp_dof', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_dof", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_tmp_pose', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_pose", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_tmp_pose', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_pose", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 1, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 384]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 384]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 1, 64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [13, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [17, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [23, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_dist', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_dist', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_q_dot_sh', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot_sh", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_q_dot_sh', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot_sh", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_time', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_time', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'q_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "q_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 's_proj', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "s_proj", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [30720]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'layer', 'd_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "d_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1216]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'layer', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [61440]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_feat', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_feat', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_scalar', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_scalar', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_time', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_time', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '3', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '3', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'q_s_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_s_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'biases', 'bias_0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "biases", "key_type": 2}, {"key": "bias_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '1')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '2')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "2", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '3')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "3", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '4')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "4", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '5')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "5", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '6')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "6", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '7')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "7", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '8')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "8", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '9')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "9", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '10')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "10", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '11')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "11", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '12')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "12", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '13')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "13", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '14')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "14", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 's_proj', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "s_proj", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [30720]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'layer', 'd_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "d_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1216]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'layer', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [61440]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_feat', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_feat', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_scalar', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_scalar', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_time', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_time', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '3', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '3', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'q_s_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_s_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'biases', 'bias_0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "biases", "key_type": 2}, {"key": "bias_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '1')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '2')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "2", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '3')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "3", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '4')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "4", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '5')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "5", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '6')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "6", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '7')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "7", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '8')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "8", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '9')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "9", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '10')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "10", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '11')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "11", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '12')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "12", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '13')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "13", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '14')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "14", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32, 192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 224]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 224]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 224]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 96]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'proj', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'proj', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 3]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'proj', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'proj', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'weight_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'weight_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'weight_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'weight_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '2', 'count')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('1', 'step')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "step", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/me-full_25000_30/_sharding b/me-full_25000_30/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..74ab7ac219f1dc2740b738c2f142a64227d66d13
--- /dev/null
+++ b/me-full_25000_30/_sharding
@@ -0,0 +1 @@
+{"MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LmF0dG5fd2VpZ2h0cw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAucHJvai5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEucHJvai5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LndlaWdodF9pcnJlcC5ydG9p":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmxuX2Rpc3RfZG9mLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmxuX2Rpc3RfZG9mLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmxuX3RtcF9kb2YuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmxuX3RtcF9kb2Yuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmxuX3RtcF9wb3NlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLmxuX3RtcF9wb3NlLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuc2RfZG90X21vZHVsYXRpb24=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuYWxpZ25fdG9fYXBwcm9hY2guMA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuYWxpZ25fdG9fYXBwcm9hY2guMQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuZmluZ2VydGlwX25vcm1hbHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuZmluZ2VydGlwX2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuaW5pdF9wcmVncmFzcF9qb2ludA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIubG9jYWxfZmluZ2VydGlwX2NvbnRhY3RfcG9zaXRpb25z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuYWxpZ25fdG9fYXBwcm9hY2guMA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuYWxpZ25fdG9fYXBwcm9hY2guMQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuZmluZ2VydGlwX25vcm1hbHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuZmluZ2VydGlwX2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuaW5pdF9wcmVncmFzcF9qb2ludA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMubG9jYWxfZmluZ2VydGlwX2NvbnRhY3RfcG9zaXRpb25z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3Rfc2hfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcud2VpZ2h0X2lycmVwLnJ0b2k=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi53ZWlnaHRfaXJyZXAucnRvaQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC53ZWlnaHRfaXJyZXAucnRvaQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi53ZWlnaHRfaXJyZXAucnRvaQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC53ZWlnaHRfaXJyZXAucnRvaQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5jb3VudA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LmF0dG5fd2VpZ2h0cw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX2Rpc3RfZG9mLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX2Rpc3RfZG9mLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX3RtcF9kb2YuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX3RtcF9kb2Yuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX3RtcF9wb3NlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX3RtcF9wb3NlLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuc2RfZG90X21vZHVsYXRpb24=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3Rfc2hfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LmF0dG5fd2VpZ2h0cw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX2Rpc3RfZG9mLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX2Rpc3RfZG9mLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX3RtcF9kb2YuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX3RtcF9kb2Yuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX3RtcF9wb3NlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX3RtcF9wb3NlLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuc2RfZG90X21vZHVsYXRpb24=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3Rfc2hfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5vcHRfc3RhdGUuMi5jb3VudA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","MS5zdGVw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}"}
\ No newline at end of file
diff --git a/me-full_25000_30/array_metadatas/process_0 b/me-full_25000_30/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..904e2ec3f8cfab2426dff227f1565193905a1667
--- /dev/null
+++ b/me-full_25000_30/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "0.decoder.attn_layers.0.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.4.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.4.kernel", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_dist_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_dist_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_dist_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_dist_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_tmp_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_tmp_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_tmp_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_tmp_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.coefficient_idx", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid.from_grid_mat", "write_shape": [6, 7, 9], "chunk_shape": [6, 7, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid.to_grid_mat", "write_shape": [6, 7, 9], "chunk_shape": [6, 7, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_1.expand_index", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_2.expand_index", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.fc_0.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.fc_0.kernel", "write_shape": [256, 384], "chunk_shape": [256, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.map.to_m_var", "write_shape": [4, 4], "chunk_shape": [4, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.4.kernel", "write_shape": [256, 384], "chunk_shape": [256, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.coefficient_idx", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid.from_grid_mat", "write_shape": [6, 7, 9], "chunk_shape": [6, 7, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid.to_grid_mat", "write_shape": [6, 7, 9], "chunk_shape": [6, 7, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_1.expand_index", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_2.expand_index", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.2.layers.1.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.2.layers.1.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.2.layers.1.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.2.layers.1.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.base_to_contact", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.dof_embedding", "write_shape": [3, 64, 9], "chunk_shape": [3, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.joint_ranges", "write_shape": [2, 2], "chunk_shape": [2, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.joint_transforms", "write_shape": [2, 6], "chunk_shape": [2, 6], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.kinematics_transforms", "write_shape": [2, 7], "chunk_shape": [2, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.1.base_to_contact", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.1.dof_embedding", "write_shape": [3, 64, 9], "chunk_shape": [3, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.1.joint_ranges", "write_shape": [2, 2], "chunk_shape": [2, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.1.joint_transforms", "write_shape": [2, 6], "chunk_shape": [2, 6], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.1.kinematics_transforms", "write_shape": [2, 7], "chunk_shape": [2, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.1.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.align_to_approach.0", "write_shape": [3, 3], "chunk_shape": [3, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.align_to_approach.1", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.base_to_contact", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.dof_embedding", "write_shape": [13, 64, 9], "chunk_shape": [13, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.fingertip_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.fingertip_normals", "write_shape": [3, 3], "chunk_shape": [3, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.init_pregrasp_joint", "write_shape": [12], "chunk_shape": [12], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.joint_ranges", "write_shape": [12, 2], "chunk_shape": [12, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.joint_transforms", "write_shape": [12, 6], "chunk_shape": [12, 6], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.kinematics_transforms", "write_shape": [12, 7], "chunk_shape": [12, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.local_fingertip_contact_positions", "write_shape": [3, 6, 3], "chunk_shape": [3, 6, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.align_to_approach.0", "write_shape": [3, 3], "chunk_shape": [3, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.align_to_approach.1", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.base_to_contact", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.dof_embedding", "write_shape": [17, 64, 9], "chunk_shape": [17, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.fingertip_idx", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.fingertip_normals", "write_shape": [4, 3], "chunk_shape": [4, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.init_pregrasp_joint", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.joint_ranges", "write_shape": [16, 2], "chunk_shape": [16, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.joint_transforms", "write_shape": [16, 6], "chunk_shape": [16, 6], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.kinematics_transforms", "write_shape": [16, 7], "chunk_shape": [16, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.local_fingertip_contact_positions", "write_shape": [4, 3, 3], "chunk_shape": [4, 3, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.4.base_to_contact", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.4.dof_embedding", "write_shape": [23, 64, 9], "chunk_shape": [23, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.4.joint_ranges", "write_shape": [22, 2], "chunk_shape": [22, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.4.joint_transforms", "write_shape": [22, 6], "chunk_shape": [22, 6], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.4.kinematics_transforms", "write_shape": [22, 7], "chunk_shape": [22, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.4.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.5.base_to_contact", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.5.dof_embedding", "write_shape": [1, 64, 9], "chunk_shape": [1, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.5.joint_ranges", "write_shape": [1, 2], "chunk_shape": [1, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.5.joint_transforms", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.5.kinematics_transforms", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.5.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_dist.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_dist.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_q_dot_sh.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_q_dot_sh.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.q_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.weight_irrep.rtoi", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.0.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.0.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.0.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.expand_index", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.s_proj.l_weights", "write_shape": [30720], "chunk_shape": [30720], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.layer.l_weights", "write_shape": [61440], "chunk_shape": [61440], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.q_s_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.weight_irrep.rtoi", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_feat.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_feat.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_query.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_query.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.weight_irrep.rtoi", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.0.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.0.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.0.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.expand_index", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.s_proj.l_weights", "write_shape": [30720], "chunk_shape": [30720], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.layer.l_weights", "write_shape": [61440], "chunk_shape": [61440], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.q_s_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.weight_irrep.rtoi", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_feat.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_feat.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_query.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_query.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.weight_irrep.rtoi", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.attn_weights", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.proj.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.proj.weight", "write_shape": [3, 64, 32], "chunk_shape": [3, 64, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.coefficient_idx", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.grid.from_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.grid.to_grid_mat", "write_shape": [8, 9, 16], "chunk_shape": [8, 9, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_1.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_2.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.map.to_m_var", "write_shape": [9, 9], "chunk_shape": [9, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.tp_ln.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.tp_ln.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_ffn.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_ffn.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_graph.balance_degree_weight", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_graph.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.proj.expand_index", "write_shape": [9], "chunk_shape": [9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.proj.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_irrep.rtoi", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_mlp.layers.2.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.4.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.4.kernel", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_dist_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_dist_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_dist_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_dist_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_tmp_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_tmp_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_tmp_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_tmp_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.fc_0.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.fc_0.kernel", "write_shape": [256, 384], "chunk_shape": [256, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.4.kernel", "write_shape": [256, 384], "chunk_shape": [256, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.2.layers.1.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.2.layers.1.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.2.layers.1.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.2.layers.1.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.0.dof_embedding", "write_shape": [3, 64, 9], "chunk_shape": [3, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.0.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.1.dof_embedding", "write_shape": [3, 64, 9], "chunk_shape": [3, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.1.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.2.dof_embedding", "write_shape": [13, 64, 9], "chunk_shape": [13, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.2.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.3.dof_embedding", "write_shape": [17, 64, 9], "chunk_shape": [17, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.3.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.4.dof_embedding", "write_shape": [23, 64, 9], "chunk_shape": [23, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.4.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.5.dof_embedding", "write_shape": [1, 64, 9], "chunk_shape": [1, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.5.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_dist.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_dist.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_q_dot_sh.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_q_dot_sh.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.q_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.0.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.s_proj.l_weights", "write_shape": [30720], "chunk_shape": [30720], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.layer.l_weights", "write_shape": [61440], "chunk_shape": [61440], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.q_s_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.0.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.s_proj.l_weights", "write_shape": [30720], "chunk_shape": [30720], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.layer.l_weights", "write_shape": [61440], "chunk_shape": [61440], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.q_s_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.attn_weights", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.proj.weight", "write_shape": [3, 64, 32], "chunk_shape": [3, 64, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.proj.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.weight_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.weight_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.weight_mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.weight_mlp.layers.2.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.4.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.4.kernel", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_dist_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_dist_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_dist_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_dist_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_tmp_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_tmp_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_tmp_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_tmp_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.fc_0.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.fc_0.kernel", "write_shape": [256, 384], "chunk_shape": [256, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.4.kernel", "write_shape": [256, 384], "chunk_shape": [256, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.2.layers.1.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.2.layers.1.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.0.scale", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.2.layers.1.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.2.layers.1.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.0.dof_embedding", "write_shape": [3, 64, 9], "chunk_shape": [3, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.0.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.1.dof_embedding", "write_shape": [3, 64, 9], "chunk_shape": [3, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.1.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.2.dof_embedding", "write_shape": [13, 64, 9], "chunk_shape": [13, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.2.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.3.dof_embedding", "write_shape": [17, 64, 9], "chunk_shape": [17, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.3.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.4.dof_embedding", "write_shape": [23, 64, 9], "chunk_shape": [23, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.4.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.5.dof_embedding", "write_shape": [1, 64, 9], "chunk_shape": [1, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.5.pose_embedding", "write_shape": [2, 64, 9], "chunk_shape": [2, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_dist.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_dist.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_q_dot_sh.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_q_dot_sh.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.q_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.0.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.s_proj.l_weights", "write_shape": [30720], "chunk_shape": [30720], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.layer.l_weights", "write_shape": [61440], "chunk_shape": [61440], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.q_s_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.0.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.s_proj.l_weights", "write_shape": [30720], "chunk_shape": [30720], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.layer.l_weights", "write_shape": [61440], "chunk_shape": [61440], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.q_dot_sh_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.q_s_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.kernel", "write_shape": [6, 32, 192], "chunk_shape": [6, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.bias", "write_shape": [6, 32], "chunk_shape": [6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.kernel", "write_shape": [192, 6, 32], "chunk_shape": [192, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.d_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 192], "chunk_shape": [256, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.sd_dot_modulation", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.attn_weights", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.so3_linear_1.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.so3_linear_2.weight", "write_shape": [3, 32, 32], "chunk_shape": [3, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.fc_0.kernel", "write_shape": [192, 224], "chunk_shape": [192, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 384], "chunk_shape": [128, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.fc_0.bias", "write_shape": [96], "chunk_shape": [96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.fc_0.kernel", "write_shape": [96, 96], "chunk_shape": [96, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 192], "chunk_shape": [128, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.tp_ln.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_ffn.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_graph.affine_weight", "write_shape": [32, 3], "chunk_shape": [32, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.proj.weight", "write_shape": [3, 64, 32], "chunk_shape": [3, 64, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.so3_linear_1.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.so3_linear_2.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.fc_0.kernel", "write_shape": [384, 448], "chunk_shape": [384, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [192, 768], "chunk_shape": [192, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.fc_0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.fc_0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [192], "chunk_shape": [192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [192, 192], "chunk_shape": [192, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [192, 384], "chunk_shape": [192, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.tp_ln.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_ffn.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_graph.affine_weight", "write_shape": [64, 3], "chunk_shape": [64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.proj.weight", "write_shape": [3, 64, 64], "chunk_shape": [3, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.weight_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.weight_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.weight_mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.weight_mlp.layers.2.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.2.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.step", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/me-full_25000_30/d/ed0c67802ee89a318ab7d7957e50af2b b/me-full_25000_30/d/ed0c67802ee89a318ab7d7957e50af2b
new file mode 100644
index 0000000000000000000000000000000000000000..0964902b8e077a2e85b79c2e5a8d3aabfb9779be
--- /dev/null
+++ b/me-full_25000_30/d/ed0c67802ee89a318ab7d7957e50af2b
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cdb8346dd9d03fb833494242bf71b825ae71c6816e47afd3475c1165d9780d36
+size 718143
diff --git a/me-full_25000_30/manifest.ocdbt b/me-full_25000_30/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..baf3c0e2e5627e9e351fc8d853f859e86bc6c12a
Binary files /dev/null and b/me-full_25000_30/manifest.ocdbt differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/0c19e7970e166aacfa715654f2f1bd51 b/me-full_25000_30/ocdbt.process_0/d/0c19e7970e166aacfa715654f2f1bd51
new file mode 100644
index 0000000000000000000000000000000000000000..fc2514277cced7a5945698b86ea58c3faef8859a
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/0c19e7970e166aacfa715654f2f1bd51 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/139cb3d7595fb4ad0d92660b2e36641d b/me-full_25000_30/ocdbt.process_0/d/139cb3d7595fb4ad0d92660b2e36641d
new file mode 100644
index 0000000000000000000000000000000000000000..05235c5e708d160be889276a03703aa89c41433f
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/139cb3d7595fb4ad0d92660b2e36641d differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/1a3e43e6db2768254dd8893337df5c8c b/me-full_25000_30/ocdbt.process_0/d/1a3e43e6db2768254dd8893337df5c8c
new file mode 100644
index 0000000000000000000000000000000000000000..ddc7f176ad1c5fc0143539f6d8ce543827a38e89
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/1a3e43e6db2768254dd8893337df5c8c differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/1bd0f148d0e6622c1580aacaf240f88d b/me-full_25000_30/ocdbt.process_0/d/1bd0f148d0e6622c1580aacaf240f88d
new file mode 100644
index 0000000000000000000000000000000000000000..7f511279cdc3a97f7855b6d927b02df055d57c44
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/1bd0f148d0e6622c1580aacaf240f88d differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/1ce129514871f69ea5f43c12d0809f78 b/me-full_25000_30/ocdbt.process_0/d/1ce129514871f69ea5f43c12d0809f78
new file mode 100644
index 0000000000000000000000000000000000000000..1866f20d3015ff204597603b1ae70c8253dd9494
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/1ce129514871f69ea5f43c12d0809f78 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/1ed05967a7b26bfbd50936dd6f2b96b1 b/me-full_25000_30/ocdbt.process_0/d/1ed05967a7b26bfbd50936dd6f2b96b1
new file mode 100644
index 0000000000000000000000000000000000000000..ce38a7e28dbfc7595b11cda0f9b7977e3352950e
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/1ed05967a7b26bfbd50936dd6f2b96b1 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/202817eb6055d2ebf5feb12cab3bf9d3 b/me-full_25000_30/ocdbt.process_0/d/202817eb6055d2ebf5feb12cab3bf9d3
new file mode 100644
index 0000000000000000000000000000000000000000..df6e9f728a699c4adc977bbda1f983c4ac3cf05e
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/202817eb6055d2ebf5feb12cab3bf9d3 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/2592ddb0c3f3499e55473b2fe737b1da b/me-full_25000_30/ocdbt.process_0/d/2592ddb0c3f3499e55473b2fe737b1da
new file mode 100644
index 0000000000000000000000000000000000000000..1d2c87c74772594fbf2e52e1e395b5926777ca9a
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/2592ddb0c3f3499e55473b2fe737b1da differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/2a55aa3bc7d7fea4e85032f9308dbc8d b/me-full_25000_30/ocdbt.process_0/d/2a55aa3bc7d7fea4e85032f9308dbc8d
new file mode 100644
index 0000000000000000000000000000000000000000..585849eb164ee7ef65b74dc09baa79ac1ae60287
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/2a55aa3bc7d7fea4e85032f9308dbc8d differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/2fdab9a68affa22afafb311936b4e917 b/me-full_25000_30/ocdbt.process_0/d/2fdab9a68affa22afafb311936b4e917
new file mode 100644
index 0000000000000000000000000000000000000000..dfd91272acf955d925c3388efea46d696dfb1efa
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/2fdab9a68affa22afafb311936b4e917 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/39a6249b85ab5b6b6a6cd7ca1b476d0b b/me-full_25000_30/ocdbt.process_0/d/39a6249b85ab5b6b6a6cd7ca1b476d0b
new file mode 100644
index 0000000000000000000000000000000000000000..a2ef52aa036444ce3e08c2be229a6d4e5e1f6961
--- /dev/null
+++ b/me-full_25000_30/ocdbt.process_0/d/39a6249b85ab5b6b6a6cd7ca1b476d0b
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:83ba1468a357d7bb8b7acff4bc376cd344e107fd1f46014910114e92476ac7a1
+size 37195776
diff --git a/me-full_25000_30/ocdbt.process_0/d/39a706121a3a25f47a5d1de3bae6d70a b/me-full_25000_30/ocdbt.process_0/d/39a706121a3a25f47a5d1de3bae6d70a
new file mode 100644
index 0000000000000000000000000000000000000000..539568cc38ed38a777b6c248f67924303be0ef55
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/39a706121a3a25f47a5d1de3bae6d70a differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/3ea3b114cc23e9e868f85e851f3b6790 b/me-full_25000_30/ocdbt.process_0/d/3ea3b114cc23e9e868f85e851f3b6790
new file mode 100644
index 0000000000000000000000000000000000000000..dc17f3d0eda6d023a71d715e4cdf7fbf76fe3c27
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/3ea3b114cc23e9e868f85e851f3b6790 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/407d70aefb2d9ac0d44311fa0773502b b/me-full_25000_30/ocdbt.process_0/d/407d70aefb2d9ac0d44311fa0773502b
new file mode 100644
index 0000000000000000000000000000000000000000..011ea8af1ad0626485789c34de5371a9e1405c04
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/407d70aefb2d9ac0d44311fa0773502b differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/438064f8912286e5c65301d98c3cde35 b/me-full_25000_30/ocdbt.process_0/d/438064f8912286e5c65301d98c3cde35
new file mode 100644
index 0000000000000000000000000000000000000000..1e3ff301e1191e90df83d1776134f8d66c8f104f
--- /dev/null
+++ b/me-full_25000_30/ocdbt.process_0/d/438064f8912286e5c65301d98c3cde35
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ad2fb8254e92fcd3033230b0dd1447da2c84aaa746c7de65f276da1c73684dab
+size 168968192
diff --git a/me-full_25000_30/ocdbt.process_0/d/45830d4b20f3ec69dee6afc813c38b8a b/me-full_25000_30/ocdbt.process_0/d/45830d4b20f3ec69dee6afc813c38b8a
new file mode 100644
index 0000000000000000000000000000000000000000..22024b8db807a122690dd6da527144bc3847b43a
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/45830d4b20f3ec69dee6afc813c38b8a differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/4b45f8a7152cd9a5e81dfdf3a19c1c08 b/me-full_25000_30/ocdbt.process_0/d/4b45f8a7152cd9a5e81dfdf3a19c1c08
new file mode 100644
index 0000000000000000000000000000000000000000..419ba26580cc2a38439cc88d72ed6e6d6638e3cc
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/4b45f8a7152cd9a5e81dfdf3a19c1c08 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/4f360d0ab9ab9b15224196afcbdcd44c b/me-full_25000_30/ocdbt.process_0/d/4f360d0ab9ab9b15224196afcbdcd44c
new file mode 100644
index 0000000000000000000000000000000000000000..db75347f00a9276e258cf0074d90285f77d45c86
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/4f360d0ab9ab9b15224196afcbdcd44c differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/57087405dcf01f8410cdc707a63d09a8 b/me-full_25000_30/ocdbt.process_0/d/57087405dcf01f8410cdc707a63d09a8
new file mode 100644
index 0000000000000000000000000000000000000000..6a80128b9bfc095d6d9359ba4b78edc0bed79200
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/57087405dcf01f8410cdc707a63d09a8 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/597e36e10180e3d8742cd5c91f6ad871 b/me-full_25000_30/ocdbt.process_0/d/597e36e10180e3d8742cd5c91f6ad871
new file mode 100644
index 0000000000000000000000000000000000000000..5adc3d50426ba97fcb86f35cba19a39375bb6d18
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/597e36e10180e3d8742cd5c91f6ad871 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/5b149e4202596f64302e04664823f55e b/me-full_25000_30/ocdbt.process_0/d/5b149e4202596f64302e04664823f55e
new file mode 100644
index 0000000000000000000000000000000000000000..9c1d33cd742766a933788c9c5bd0dcc5e4ce7323
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/5b149e4202596f64302e04664823f55e differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/5fde28dc845becb91a81931d5fecf495 b/me-full_25000_30/ocdbt.process_0/d/5fde28dc845becb91a81931d5fecf495
new file mode 100644
index 0000000000000000000000000000000000000000..8a5c90be7b7d4942b0e028af8f4e08e13aa9579f
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/5fde28dc845becb91a81931d5fecf495 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/652a52de61856b5a92e387057eaef9b0 b/me-full_25000_30/ocdbt.process_0/d/652a52de61856b5a92e387057eaef9b0
new file mode 100644
index 0000000000000000000000000000000000000000..3bf204c2d65d66ce3f497a947e5430b53cf7bbe6
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/652a52de61856b5a92e387057eaef9b0 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/6ace53f61ab846c2b95b0e08a9013d82 b/me-full_25000_30/ocdbt.process_0/d/6ace53f61ab846c2b95b0e08a9013d82
new file mode 100644
index 0000000000000000000000000000000000000000..ad03da7050bec695320d578f94163052a9f2bba2
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/6ace53f61ab846c2b95b0e08a9013d82 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/6c6b5f9ad89bd45eb26cc27a87e3847f b/me-full_25000_30/ocdbt.process_0/d/6c6b5f9ad89bd45eb26cc27a87e3847f
new file mode 100644
index 0000000000000000000000000000000000000000..9b968feb8dae8d5c0fd40b25af6fda32c244f60a
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/6c6b5f9ad89bd45eb26cc27a87e3847f differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/716c59e8e25a0a1f1d7e028d02a32618 b/me-full_25000_30/ocdbt.process_0/d/716c59e8e25a0a1f1d7e028d02a32618
new file mode 100644
index 0000000000000000000000000000000000000000..f9be7e9fef13c5850e79b437209e45733f7d7afd
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/716c59e8e25a0a1f1d7e028d02a32618 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/734d35072c33ecfc8ee2c4889a42651e b/me-full_25000_30/ocdbt.process_0/d/734d35072c33ecfc8ee2c4889a42651e
new file mode 100644
index 0000000000000000000000000000000000000000..a279063880e96eae768619ebdb175f271db1fb76
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/734d35072c33ecfc8ee2c4889a42651e differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/780b89073f009dbbb3fbb685a4de10ef b/me-full_25000_30/ocdbt.process_0/d/780b89073f009dbbb3fbb685a4de10ef
new file mode 100644
index 0000000000000000000000000000000000000000..2d88eb137bc7b1713550dea0d88b5a5f944a037d
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/780b89073f009dbbb3fbb685a4de10ef differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/7d42c16e0858d454abe06facc53ea83a b/me-full_25000_30/ocdbt.process_0/d/7d42c16e0858d454abe06facc53ea83a
new file mode 100644
index 0000000000000000000000000000000000000000..3240d6ed726877c798cb2cf2234b4542a0648afd
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/7d42c16e0858d454abe06facc53ea83a differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/8311c8e07bef5ae84b978a6bc04c8b53 b/me-full_25000_30/ocdbt.process_0/d/8311c8e07bef5ae84b978a6bc04c8b53
new file mode 100644
index 0000000000000000000000000000000000000000..2cb18e4180b6bca51d918f13c29748daec3b3e37
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/8311c8e07bef5ae84b978a6bc04c8b53 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/8da9633f14b6edd22ba730f20d2ac9e5 b/me-full_25000_30/ocdbt.process_0/d/8da9633f14b6edd22ba730f20d2ac9e5
new file mode 100644
index 0000000000000000000000000000000000000000..3e9b700f099e40c49ffa16d55dfb9bc88f692c6b
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/8da9633f14b6edd22ba730f20d2ac9e5 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/8e65eec0f9a2378c5ba09951cefa45ab b/me-full_25000_30/ocdbt.process_0/d/8e65eec0f9a2378c5ba09951cefa45ab
new file mode 100644
index 0000000000000000000000000000000000000000..b464a5b87aa59faddcdbc39f0d52f64a4979161d
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/8e65eec0f9a2378c5ba09951cefa45ab differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/a0ad73b7275e26df75a83c65bd102022 b/me-full_25000_30/ocdbt.process_0/d/a0ad73b7275e26df75a83c65bd102022
new file mode 100644
index 0000000000000000000000000000000000000000..7a7c2449d7348d6695d8943838e872bbdfb94063
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/a0ad73b7275e26df75a83c65bd102022 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/b11c7ec964952e225499f6da3bf0dd13 b/me-full_25000_30/ocdbt.process_0/d/b11c7ec964952e225499f6da3bf0dd13
new file mode 100644
index 0000000000000000000000000000000000000000..eabf3d797b20350d829ce9e6304d57d749a85c96
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/b11c7ec964952e225499f6da3bf0dd13 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/b6d66a5a1b16d08a3b1e13bd95cd135e b/me-full_25000_30/ocdbt.process_0/d/b6d66a5a1b16d08a3b1e13bd95cd135e
new file mode 100644
index 0000000000000000000000000000000000000000..2a3696cf019934da69b5429095799f9e454a6d5e
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/b6d66a5a1b16d08a3b1e13bd95cd135e differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/b8465c24058b8f0dd5d79ecbc241a840 b/me-full_25000_30/ocdbt.process_0/d/b8465c24058b8f0dd5d79ecbc241a840
new file mode 100644
index 0000000000000000000000000000000000000000..5d0029b1cc3783b6f8255bd75fab41d76b5f0c37
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/b8465c24058b8f0dd5d79ecbc241a840 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/b88844ecb87be367b703b466e876d71a b/me-full_25000_30/ocdbt.process_0/d/b88844ecb87be367b703b466e876d71a
new file mode 100644
index 0000000000000000000000000000000000000000..b33e17bb0118e34f20f6c290919f58e39db21d1a
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/b88844ecb87be367b703b466e876d71a differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/c96f93e7a3263366a8d189236b32238a b/me-full_25000_30/ocdbt.process_0/d/c96f93e7a3263366a8d189236b32238a
new file mode 100644
index 0000000000000000000000000000000000000000..a55f0a9d9ea3518329a82e1b8fda78f44ddd3aff
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/c96f93e7a3263366a8d189236b32238a differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/d213560bd48ed855440c462191845f39 b/me-full_25000_30/ocdbt.process_0/d/d213560bd48ed855440c462191845f39
new file mode 100644
index 0000000000000000000000000000000000000000..3bb900c388b663a06dbd2b54800989a5576c4fff
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/d213560bd48ed855440c462191845f39 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/d3218fa39ee5f86a0c1cf8bfdc8db625 b/me-full_25000_30/ocdbt.process_0/d/d3218fa39ee5f86a0c1cf8bfdc8db625
new file mode 100644
index 0000000000000000000000000000000000000000..f6555fb564769d5cb36160b7396f8534826cd75e
--- /dev/null
+++ b/me-full_25000_30/ocdbt.process_0/d/d3218fa39ee5f86a0c1cf8bfdc8db625
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d485ec6f58407821acd28aa70959094faff0e416d10d9eab89be822ce7b484fa
+size 1270964
diff --git a/me-full_25000_30/ocdbt.process_0/d/da0a94fd479b58bb128ef4cf1574cb4d b/me-full_25000_30/ocdbt.process_0/d/da0a94fd479b58bb128ef4cf1574cb4d
new file mode 100644
index 0000000000000000000000000000000000000000..e9ad46c5b18d8e1aa0898b0ea692fdf2f05cbb5c
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/da0a94fd479b58bb128ef4cf1574cb4d differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/f14ddd9b34ef7b08660bff428de4ee42 b/me-full_25000_30/ocdbt.process_0/d/f14ddd9b34ef7b08660bff428de4ee42
new file mode 100644
index 0000000000000000000000000000000000000000..313f07ec287a9df6bb5368c8032a9a1e011651d6
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/f14ddd9b34ef7b08660bff428de4ee42 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/f3a52b4a9f6f653606de42f0945c8db2 b/me-full_25000_30/ocdbt.process_0/d/f3a52b4a9f6f653606de42f0945c8db2
new file mode 100644
index 0000000000000000000000000000000000000000..8e79d62edcc2b2e565c28ca8b06831d4129b81d6
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/f3a52b4a9f6f653606de42f0945c8db2 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/f437b0e4cead6bdbe0489ccc53c81891 b/me-full_25000_30/ocdbt.process_0/d/f437b0e4cead6bdbe0489ccc53c81891
new file mode 100644
index 0000000000000000000000000000000000000000..95c87351a57e543e256ab5ea9e5b4d7c62f81b28
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/f437b0e4cead6bdbe0489ccc53c81891 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/f64af71e2a1fd20467a13ddc3fb3b2f1 b/me-full_25000_30/ocdbt.process_0/d/f64af71e2a1fd20467a13ddc3fb3b2f1
new file mode 100644
index 0000000000000000000000000000000000000000..9a1b1e6c6568bdf8b171503d7cf0ca13a7f6e0d2
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/f64af71e2a1fd20467a13ddc3fb3b2f1 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/f88cc61ae1a15ce7ddc96246fa8a6d27 b/me-full_25000_30/ocdbt.process_0/d/f88cc61ae1a15ce7ddc96246fa8a6d27
new file mode 100644
index 0000000000000000000000000000000000000000..088ad61499fd20b6e4d8b09158c199eebe312885
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/f88cc61ae1a15ce7ddc96246fa8a6d27 differ
diff --git a/me-full_25000_30/ocdbt.process_0/d/fae290a2c5a0f77689db0ec3b37b9256 b/me-full_25000_30/ocdbt.process_0/d/fae290a2c5a0f77689db0ec3b37b9256
new file mode 100644
index 0000000000000000000000000000000000000000..70d8f7c854df7d6121fc48a511f32f91dc84a0f2
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/d/fae290a2c5a0f77689db0ec3b37b9256 differ
diff --git a/me-full_25000_30/ocdbt.process_0/manifest.ocdbt b/me-full_25000_30/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..dc9e74d130e58222c71d31c1e7331cbec6f58cb2
Binary files /dev/null and b/me-full_25000_30/ocdbt.process_0/manifest.ocdbt differ