diff --git a/.gitattributes b/.gitattributes
index e2e8232b9f4daf7b3526983f0536577f21d17bbd..b4430ba4c1fdd92f5dc48b3a1ae2d12502852502 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -66,3 +66,8 @@ me-full_25000_120/ocdbt.process_0/d/3b2682464c793b995e3ae711a58a340a filter=lfs
 me-full_25000_120/ocdbt.process_0/d/4d41c06c7629cf3abcd6206edbf5d41b filter=lfs diff=lfs merge=lfs -text
 me-full_25000_120/ocdbt.process_0/d/786ed4041c8ed2514103f5daa63b5a9d filter=lfs diff=lfs merge=lfs -text
 me-full_25000_120/ocdbt.process_0/d/edd67c80bdaf215c08a98f9635ded706 filter=lfs diff=lfs merge=lfs -text
+me-seed-anneal_25000_4/d/6df16f8850d87665698b0829137fd441 filter=lfs diff=lfs merge=lfs -text
+me-seed-anneal_25000_4/ocdbt.process_0/d/3d0f0c9004bd81f01d77a5f50e8e3329 filter=lfs diff=lfs merge=lfs -text
+me-seed-anneal_25000_4/ocdbt.process_0/d/60aca8647e3d5929295e5ceaf9c045c3 filter=lfs diff=lfs merge=lfs -text
+me-seed-anneal_25000_4/ocdbt.process_0/d/a0b5d0bb7f2ad0b02d8a64669ff910b8 filter=lfs diff=lfs merge=lfs -text
+me-seed-anneal_25000_4/ocdbt.process_0/d/e6f4a07d5a6435f3c5dab930d72e9ae6 filter=lfs diff=lfs merge=lfs -text
diff --git a/me-seed-anneal_25000_4/_CHECKPOINT_METADATA b/me-seed-anneal_25000_4/_CHECKPOINT_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..85e65b3a92d94bd60a3acd7531a8e57f9dc9d880
--- /dev/null
+++ b/me-seed-anneal_25000_4/_CHECKPOINT_METADATA
@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786869517801581596, "commit_timestamp_nsecs": 1786869526148822206, "custom_metadata": {}}
\ No newline at end of file
diff --git a/me-seed-anneal_25000_4/_METADATA b/me-seed-anneal_25000_4/_METADATA
new file mode 100644
index 0000000000000000000000000000000000000000..34e12fdaa3372cfbfc072b292f31e35a9fd1d8b5
--- /dev/null
+++ b/me-seed-anneal_25000_4/_METADATA
@@ -0,0 +1 @@
+{"tree_metadata": {"('0', 'decoder', 'attn_layers', '0', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '0', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('0', 'decoder', 'attn_layers', '0', '1', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('0', 'decoder', 'attn_layers', '0', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '0', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'attn_layers', '1', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '1', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('0', 'decoder', 'attn_layers', '1', '1', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('0', 'decoder', 'attn_layers', '1', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'attn_layers', '1', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'dof_gl', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'decoder', 'dof_gl', 'd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'decoder', 'dof_gl', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'decoder', 'dof_gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('0', 'decoder', 'dof_gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'decoder', 'dof_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'dof_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'dof_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'dof_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('0', 'decoder', 'dof_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'decoder', 'ln_dist_dof', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_dof", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_dist_dof', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_dof", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_dist_pose', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_pose", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_dist_pose', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_pose", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_tmp_dof', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_dof", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_tmp_dof', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_dof", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_tmp_pose', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_pose", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'ln_tmp_pose', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_pose", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'decoder', 'pos_head', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 7, 9]}}, "('0', 'decoder', 'pos_head', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 7, 9]}}, "('0', 'decoder', 'pos_head', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'pos_head', 'grid_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'grid_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'pos_head', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pos_head', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'decoder', 'pos_head', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 64]}}, "('0', 'decoder', 'pos_head', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('0', 'decoder', 'pos_head', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'decoder', 'pos_head', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 1, 64]}}, "('0', 'decoder', 'pose_gl', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'decoder', 'pose_gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pose_gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pose_gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pose_gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 128]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 4]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 128]}}, "('0', 'decoder', 'pose_gl', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'decoder', 'rot_head', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'decoder', 'rot_head', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 7, 9]}}, "('0', 'decoder', 'rot_head', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 7, 9]}}, "('0', 'decoder', 'rot_head', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'rot_head', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'rot_head', 'grid_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'rot_head', 'grid_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'decoder', 'rot_head', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'decoder', 'rot_head', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'decoder', 'rot_head', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 64]}}, "('0', 'decoder', 'rot_head', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('0', 'decoder', 'rot_head', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'decoder', 'rot_head', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 1, 64]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32, 192]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '1', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'attn_layers', '0', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32, 192]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '1', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'attn_layers', '1', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'kinematic_encoding', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('0', 'kinematic_encoding', 'gl', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('0', 'kinematic_encoding', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'kinematic_encoding', 'gl', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'kinematic_encoding', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('0', 'kinematic_encoding', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'base_to_contact')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "base_to_contact", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'dof_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'joint_ranges')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "joint_ranges", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'joint_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "joint_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'kinematics_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kinematics_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'pose_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'base_to_contact')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "base_to_contact", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'dof_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 9]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'joint_ranges')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "joint_ranges", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'joint_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "joint_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'kinematics_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kinematics_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'pose_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'align_to_approach', '0')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "align_to_approach", "key_type": 2}, {"key": "0", "key_type": 1}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'align_to_approach', '1')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "align_to_approach", "key_type": 2}, {"key": "1", "key_type": 1}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'base_to_contact')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "base_to_contact", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'dof_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [13, 64, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'fingertip_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fingertip_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'fingertip_normals')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fingertip_normals", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'init_pregrasp_joint')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "init_pregrasp_joint", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'joint_ranges')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "joint_ranges", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'joint_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "joint_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 6]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'kinematics_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kinematics_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'local_fingertip_contact_positions')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "local_fingertip_contact_positions", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 6, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'pose_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'align_to_approach', '0')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "align_to_approach", "key_type": 2}, {"key": "0", "key_type": 1}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'align_to_approach', '1')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "align_to_approach", "key_type": 2}, {"key": "1", "key_type": 1}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'base_to_contact')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "base_to_contact", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'dof_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [17, 64, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'fingertip_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "fingertip_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'fingertip_normals')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "fingertip_normals", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 1]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'init_pregrasp_joint')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "init_pregrasp_joint", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'joint_ranges')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "joint_ranges", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'joint_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "joint_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'kinematics_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kinematics_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'local_fingertip_contact_positions')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "local_fingertip_contact_positions", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 1, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'pose_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'base_to_contact')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "base_to_contact", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'dof_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [23, 64, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'joint_ranges')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "joint_ranges", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [22, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'joint_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "joint_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [22, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'kinematics_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kinematics_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [22, 7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'pose_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'base_to_contact')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "base_to_contact", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [7]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'dof_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 3]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'joint_ranges')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "joint_ranges", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 2]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'joint_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "joint_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'kinematics_transforms')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "kinematics_transforms", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('0', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'pose_embedding')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 3]}}, "('0', 'kinematic_encoding', 'ln_dist', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_dist', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_q_dot_sh', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot_sh", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_q_dot_sh', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot_sh", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_time', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_time', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'kinematic_encoding', 'mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('0', 'kinematic_encoding', 'q_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "q_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'kinematic_encoding', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'kinematic_encoding', 'weight_irrep', 'rtoi')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "weight_irrep", "key_type": 2}, {"key": "rtoi", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 's_proj', 'l_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "s_proj", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [10240]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'layer', 'd_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "d_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1216]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'layer', 'l_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20480]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_feat', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_feat', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_scalar', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_scalar', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_time', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_time', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '3', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '3', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'q_s_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_s_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'weight_irrep', 'rtoi')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "weight_irrep", "key_type": 2}, {"key": "rtoi", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'biases', 'bias_0')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "biases", "key_type": 2}, {"key": "bias_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '0')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '1')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '2')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "2", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '3')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "3", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '4')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "4", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '5')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "5", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '6')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "6", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '7')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "7", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '8')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "8", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '9')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "9", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '10')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "10", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '11')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "11", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '12')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "12", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '13')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "13", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '14')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "14", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32, 192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32, 192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'weight_irrep', 'rtoi')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "weight_irrep", "key_type": 2}, {"key": "rtoi", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 's_proj', 'l_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "s_proj", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [10240]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'layer', 'd_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "d_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1216]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'layer', 'l_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20480]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_feat', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_feat', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_scalar', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_scalar', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_time', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_time', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '3', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '3', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'q_s_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_s_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'weight_irrep', 'rtoi')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "weight_irrep", "key_type": 2}, {"key": "rtoi", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'biases', 'bias_0')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "biases", "key_type": 2}, {"key": "bias_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '0')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '1')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '2')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "2", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '3')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "3", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '4')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "4", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '5')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "5", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '6')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "6", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '7')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "7", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '8')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "8", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '9')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "9", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '10')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "10", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '11')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "11", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '12')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "12", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '13')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "13", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '14')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "14", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32, 192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'key', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'key', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'out', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'out', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32, 192]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'query', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'query', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'value', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'value', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'weight_irrep', 'rtoi')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "weight_irrep", "key_type": 2}, {"key": "rtoi", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 224]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 96]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 64]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 224]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 96]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 64]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 224]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 96]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 64]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'pool_blocks', '0', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '0', 'proj', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '0', 'proj', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '0', 'proj', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 32]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'coefficient_idx')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "coefficient_idx", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid', 'from_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "from_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid', 'to_grid_mat')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid", "key_type": 2}, {"key": "to_grid_mat", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 3, 16]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'map', 'to_m_var')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "map", "key_type": 2}, {"key": "to_m_var", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 9]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_ffn', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_ffn', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_graph', 'balance_degree_weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "balance_degree_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 3]}}, "('0', 'unet', 'pool_blocks', '1', 'ln_graph', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '1', 'proj', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'pool_blocks', '1', 'proj', 'expand_index')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "expand_index", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3]}}, "('0', 'unet', 'pool_blocks', '1', 'proj', 'weight')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('0', 'unet', 'seed_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "seed_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'seed_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "seed_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'seed_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "seed_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('0', 'unet', 'seed_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "seed_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 32]}}, "('0', 'unet', 'weight_irrep', 'rtoi')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_irrep", "key_type": 2}, {"key": "rtoi", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 9]}}, "('0', 'unet', 'weight_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('0', 'unet', 'weight_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('0', 'unet', 'weight_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('0', 'unet', 'weight_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "0", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'count')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '0', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'attn_layers', '1', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'dof_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_dist_dof', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_dof", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_dist_dof', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_dof", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_dist_pose', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_pose", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_dist_pose', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_pose", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_tmp_dof', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_dof", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_tmp_dof', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_dof", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_tmp_pose', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_pose", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'ln_tmp_pose', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_pose", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pos_head', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 1, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 128]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 128]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'pose_gl', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('1', 'opt_state', '0', 'mu', 'decoder', 'rot_head', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 1, 64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32, 192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '0', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32, 192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'attn_layers', '1', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 3]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 9]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 3]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [13, 64, 3]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 3]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [17, 64, 3]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 3]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [23, 64, 3]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 3]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 3]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 3]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_dist', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_dist', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_q_dot_sh', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot_sh", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_q_dot_sh', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot_sh", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_time', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_time', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'q_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "q_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'kinematic_encoding', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 's_proj', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "s_proj", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [10240]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'layer', 'd_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "d_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1216]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'layer', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20480]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_feat', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_feat', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_scalar', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_scalar', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_time', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_time', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '3', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '3', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'q_s_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_s_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'biases', 'bias_0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "biases", "key_type": 2}, {"key": "bias_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '1')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '2')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "2", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '3')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "3", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '4')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "4", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '5')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "5", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '6')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "6", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '7')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "7", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '8')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "8", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '9')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "9", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '10')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "10", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '11')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "11", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '12')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "12", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '13')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "13", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '14')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "14", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32, 192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32, 192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 's_proj', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "s_proj", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [10240]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'layer', 'd_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "d_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1216]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'layer', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20480]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_feat', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_feat', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_scalar', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_scalar', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_time', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_time', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '3', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '3', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'q_s_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_s_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'biases', 'bias_0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "biases", "key_type": 2}, {"key": "bias_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '1')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '2')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "2", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '3')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "3", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '4')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "4", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '5')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "5", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '6')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "6", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '7')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "7", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '8')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "8", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '9')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "9", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '10')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "10", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '11')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "11", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '12')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "12", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '13')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "13", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '14')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "14", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32, 192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32, 192]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 224]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 96]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 224]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 96]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 224]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 96]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'proj', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '0', 'proj', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'proj', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'pool_blocks', '1', 'proj', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'seed_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "seed_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'seed_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "seed_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'seed_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "seed_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'seed_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "seed_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 32]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'weight_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'weight_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'weight_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'mu', 'unet', 'weight_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "mu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '0', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 4, 16]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'attn_layers', '1', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'dof_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "dof_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_dist_dof', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_dof", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_dist_dof', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_dof", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_dist_pose', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_pose", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_dist_pose', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_dist_pose", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_tmp_dof', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_dof", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_tmp_dof', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_dof", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_tmp_pose', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_pose", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'ln_tmp_pose', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "ln_tmp_pose", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'grid_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'scalar_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pos_head', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pos_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 1, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 128]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 128]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'pose_gl', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "pose_gl", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'grid_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'scalar_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1]}}, "('1', 'opt_state', '0', 'nu', 'decoder', 'rot_head', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "rot_head", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 1, 64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32, 192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '0', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32, 192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '2', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'attn_layers', '1', '2', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '0', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 3]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 9]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '1', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 3]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [13, 64, 3]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '2', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 3]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [17, 64, 3]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '3', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 3]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [23, 64, 3]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '4', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 3]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'dof_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "dof_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 3]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'kinematic_net', 'kins', '5', 'pose_embedding')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "kinematic_net", "key_type": 2}, {"key": "kins", "key_type": 2}, {"key": "5", "key_type": 2}, {"key": "pose_embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 64, 3]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_dist', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_dist', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_q_dot_sh', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot_sh", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_q_dot_sh', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_q_dot_sh", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_time', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_time', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'q_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "q_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'kinematic_encoding', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "kinematic_encoding", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '0', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_scalar_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'film', 'time_vec_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '0', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '1', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 'per_scale_ln', '2', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '1', 's_proj', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "s_proj", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [10240]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'layer', 'd_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "d_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1216]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'layer', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20480]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_feat', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_q_feat', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_scalar', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_scalar', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_time', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'ln_time', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mha', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '3', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'mlp', 'layers', '3', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '2', 'q_s_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_s_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_feat', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'ln_query', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'biases', 'bias_0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "biases", "key_type": 2}, {"key": "bias_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '1')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '2')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "2", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '3')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "3", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '4')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "4", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '5')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "5", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '6')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "6", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '7')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "7", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '8')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "8", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '9')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "9", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '10')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "10", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '11')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "11", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '12')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "12", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '13')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "13", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '3', 'tp', 'weights', '14')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "14", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32, 192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '0', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32, 192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'attn_layers', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '0', '4', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '0', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_scalar_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_scalar_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'film', 'time_vec_scales', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "film", "key_type": 2}, {"key": "time_vec_scales", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '0', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '1', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 'per_scale_ln', '2', '1', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "per_scale_ln", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '1', 's_proj', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "s_proj", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [10240]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'layer', 'd_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "d_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1216]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'layer', 'l_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "layer", "key_type": 2}, {"key": "l_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [20480]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_feat', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_q_feat', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_q_feat", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_scalar', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_scalar', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_scalar", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_time', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'ln_time', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_time", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mha', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mha", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '1', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '3', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'mlp', 'layers', '3', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'q_dot_sh_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_dot_sh_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '2', 'q_s_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "q_s_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_feat', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_feat", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'ln_query', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "ln_query", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'biases', 'bias_0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "biases", "key_type": 2}, {"key": "bias_0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '0')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "0", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '1')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "1", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '2')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "2", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '3')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "3", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '4')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "4", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '5')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "5", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '6')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "6", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '7')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "7", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '8')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "8", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '9')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "9", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '10')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "10", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '11')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "11", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '12')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "12", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '13')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "13", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '3', 'tp', 'weights', '14')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "3", "key_type": 2}, {"key": "tp", "key_type": 2}, {"key": "weights", "key_type": 2}, {"key": "14", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32, 192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '0', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'key', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'key', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'out', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'out', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32, 192]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'query', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'query', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'value', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'attn_layers', '1', 'value', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "attn_layers", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 6, 32]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "d_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_sd_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'ln_sd_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "ln_sd_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '6', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'msg_mlp', 'layers', '6', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "msg_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "6", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'gl', 'sd_dot_modulation')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "gl", "key_type": 2}, {"key": "sd_dot_modulation", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_dist_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_dist_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_dist_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_time_graph', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'multi_scale_feature_query', 'readout_blocks', '1', '4', 'ln_time_graph', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "multi_scale_feature_query", "key_type": 2}, {"key": "readout_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "ln_time_graph", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 224]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 96]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 224]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 96]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '0', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '1', '2', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'mp_blocks', '2', '2', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "mp_blocks", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 32, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [224]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 224]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 96]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'proj', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '0', 'proj', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'fnn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fnn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'attn_weights')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "attn_weights", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'grid_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "grid_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'scalar_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "scalar_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_1', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn', 'so3_linear_2', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "so3_linear_2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ffn_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ffn_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ln_d_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ln_d_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_d_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ln_s_dot', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'ln_s_dot', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "ln_s_dot", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 448]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 768]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp1', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp1", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'fc_0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'fc_0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "fc_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 192]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'mlp', 'layers', '4', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "4", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 384]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'so2_m_conv', '0', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'so2_tp2', 'so2_m_conv', '1', 'fc', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "so2_tp2", "key_type": 2}, {"key": "so2_m_conv", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "fc", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'graph', 'tp_ln', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "graph", "key_type": 2}, {"key": "tp_ln", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'ln_attr', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'ln_attr', '0', 'scale')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_attr", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'ln_ffn', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'ln_ffn', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_ffn", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'ln_graph', 'affine_bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'ln_graph', 'affine_weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ln_graph", "key_type": 2}, {"key": "affine_weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'proj', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'pool_blocks', '1', 'proj', 'weight')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "pool_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "proj", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'seed_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "seed_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'seed_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "seed_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'seed_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "seed_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'seed_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "seed_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 32]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'weight_mlp', 'layers', '0', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'weight_mlp', 'layers', '0', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 64]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'weight_mlp', 'layers', '2', 'bias')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('1', 'opt_state', '0', 'nu', 'unet', 'weight_mlp', 'layers', '2', 'kernel')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "nu", "key_type": 2}, {"key": "unet", "key_type": 2}, {"key": "weight_mlp", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 128]}}, "('1', 'opt_state', '2', 'count')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "2", "key_type": 2}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}, "('1', 'step')": {"key_metadata": [{"key": "1", "key_type": 2}, {"key": "step", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": []}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file
diff --git a/me-seed-anneal_25000_4/_sharding b/me-seed-anneal_25000_4/_sharding
new file mode 100644
index 0000000000000000000000000000000000000000..a41c54a72b12bc2f8c6f57135c83c260b8a7b589
--- /dev/null
+++ b/me-seed-anneal_25000_4/_sharding
@@ -0,0 +1 @@
+{"MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LmF0dG5fd2VpZ2h0cw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjAucHJvai5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYmFsYW5jZV9kZWdyZWVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnBvb2xfYmxvY2tzLjEucHJvai5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnNlZWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnNlZWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnNlZWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LnNlZWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LndlaWdodF9pcnJlcC5ydG9p":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmxuX2Rpc3RfZG9mLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmxuX2Rpc3RfZG9mLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmxuX3RtcF9kb2YuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmxuX3RtcF9kb2Yuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmxuX3RtcF9wb3NlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLmxuX3RtcF9wb3NlLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tYXAudG9fbV92YXI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmNvZWZmaWNpZW50X2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWQuZnJvbV9ncmlkX21hdA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWQudG9fZ3JpZF9tYXQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuc2RfZG90X21vZHVsYXRpb24=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcuZXhwYW5kX2luZGV4":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuYWxpZ25fdG9fYXBwcm9hY2guMA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuYWxpZ25fdG9fYXBwcm9hY2guMQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuZmluZ2VydGlwX25vcm1hbHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuZmluZ2VydGlwX2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuaW5pdF9wcmVncmFzcF9qb2ludA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIubG9jYWxfZmluZ2VydGlwX2NvbnRhY3RfcG9zaXRpb25z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuYWxpZ25fdG9fYXBwcm9hY2guMA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuYWxpZ25fdG9fYXBwcm9hY2guMQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuZmluZ2VydGlwX25vcm1hbHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuZmluZ2VydGlwX2lkeA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuaW5pdF9wcmVncmFzcF9qb2ludA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMubG9jYWxfZmluZ2VydGlwX2NvbnRhY3RfcG9zaXRpb25z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuYmFzZV90b19jb250YWN0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUua2luZW1hdGljc190cmFuc2Zvcm1z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuam9pbnRfcmFuZ2Vz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuam9pbnRfdHJhbnNmb3Jtcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3Rfc2hfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5raW5lbWF0aWNfZW5jb2Rpbmcud2VpZ2h0X2lycmVwLnJ0b2k=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi53ZWlnaHRfaXJyZXAucnRvaQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC53ZWlnaHRfaXJyZXAucnRvaQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi53ZWlnaHRfaXJyZXAucnRvaQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmJhbGFuY2VfZGVncmVlX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmV4cGFuZF9pbmRleA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC53ZWlnaHRfaXJyZXAucnRvaQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5iYWxhbmNlX2RlZ3JlZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MC5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5leHBhbmRfaW5kZXg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5jb3VudA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LmF0dG5fd2VpZ2h0cw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjAucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnBvb2xfYmxvY2tzLjEucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnNlZWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnNlZWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnNlZWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LnNlZWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX2Rpc3RfZG9mLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX2Rpc3RfZG9mLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX3RtcF9kb2YuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX3RtcF9kb2Yuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX3RtcF9wb3NlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLmxuX3RtcF9wb3NlLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuc2RfZG90X21vZHVsYXRpb24=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3Rfc2hfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5tdS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4wLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4xLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0Lm1wX2Jsb2Nrcy4yLjIubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LmF0dG5fd2VpZ2h0cw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjAucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguZmZuX2xuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMS5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGguc28yX3RwMi5zbzJfbV9jb252LjEuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZ3JhcGgudHBfbG4uYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEuZm5uLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fYXR0ci4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX2JpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZ3JhcGguYWZmaW5lX3dlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEubG5fZmZuLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEucHJvai53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnBvb2xfYmxvY2tzLjEucHJvai5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnNlZWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnNlZWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnNlZWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LnNlZWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS51bmV0LndlaWdodF9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjAuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS52YWx1ZS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5rZXkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5vdXQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMS5xdWVyeS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmF0dG5fbGF5ZXJzLjEuMi5sYXllcnMuMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmRvZl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX2Rpc3RfZG9mLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX2Rpc3RfZG9mLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX2Rpc3RfcG9zZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX3RtcF9kb2YuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX3RtcF9kb2Yuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX3RtcF9wb3NlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLmxuX3RtcF9wb3NlLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc19oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5mY18wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnBvc2VfZ2wuc28yX3RwMS5zbzJfbV9jb252LjAuZmMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLmdyaWRfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNjYWxhcl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5kZWNvZGVyLnJvdF9oZWFkLnNvM19saW5lYXJfMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMC4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4wLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuYXR0bl9sYXllcnMuMS4yLmxheWVycy4xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjAua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjIua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuYXR0bl9tbHAubGF5ZXJzLjQua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc19kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubG5fc2RfZG90LnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wubXNnX21scC5sYXllcnMuNi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcuZ2wuc2RfZG90X21vZHVsYXRpb24=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjAucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjEucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjIucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjMucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjQucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUuZG9mX2VtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2Rpbmcua2luZW1hdGljX25ldC5raW5zLjUucG9zZV9lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fZGlzdF9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fcV9kb3Rfc2guc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZS5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubG5fdGltZV9ncmFwaC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcubWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5raW5lbWF0aWNfZW5jb2RpbmcucV9kb3Rfc2hfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjAuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfc2NhbGFyX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5maWxtLnRpbWVfdmVjX3NjYWxlcy5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMC4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMS4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5wZXJfc2NhbGVfbG4uMi4xLmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMS5zX3Byb2oubF93ZWlnaHRz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5kX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sYXllci5sX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl90aW1lLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9xX2ZlYXQuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5sbl9zY2FsYXIuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEua2V5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEub3V0Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEucXVlcnkua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5taGEudmFsdWUua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjEua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5tbHAubGF5ZXJzLjMua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX2RvdF9zaF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMi5xX3NfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjA=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjE=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEx":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEy":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjEz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjI=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjY=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjg=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC53ZWlnaHRzLjk=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy50cC5iaWFzZXMuYmlhc18w":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV93ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9mZWF0LmFmZmluZV9iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuMy5sbl9xdWVyeS5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4wLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLm91dC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLmtleS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnF1ZXJ5Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5hdHRuX2xheWVycy4xLnZhbHVlLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuMi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX21scC5sYXllcnMuNC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5hdHRuX3dlaWdodHM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5kX2RvdF9tb2R1bGF0aW9u":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9kX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zX2RvdC5zY2FsZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3QuYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5sbl9zZF9kb3Quc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy40Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42LmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy42Lmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5tc2dfbWxwLmxheWVycy4yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5nbC5zZF9kb3RfbW9kdWxhdGlvbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl90aW1lX2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9kaXN0X2dyYXBoLnNjYWxl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMC5udS5tdWx0aV9zY2FsZV9mZWF0dXJlX3F1ZXJ5LnJlYWRvdXRfYmxvY2tzLjEuNC5sbl9ncmFwaC5hZmZpbmVfd2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5vcHRfc3RhdGUuMi5jb3VudA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}","MS5zdGVw":"{\"sharding_type\": \"NamedSharding\", \"shape\": [3], \"axis_names\": [\"devices\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}, {\"id\": 1}, {\"id\": 2}]}}"}
\ No newline at end of file
diff --git a/me-seed-anneal_25000_4/array_metadatas/process_0 b/me-seed-anneal_25000_4/array_metadatas/process_0
new file mode 100644
index 0000000000000000000000000000000000000000..cab9b2c86ee1e3c5eddae4d89c33ae2629719b66
--- /dev/null
+++ b/me-seed-anneal_25000_4/array_metadatas/process_0
@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "0.decoder.attn_layers.0.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.0.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.attn_layers.1.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.d_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_gl.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.4.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.dof_mlp.layers.4.kernel", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_dist_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_dist_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_dist_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_dist_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_tmp_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_tmp_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_tmp_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.ln_tmp_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.coefficient_idx", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid.from_grid_mat", "write_shape": [2, 7, 9], "chunk_shape": [2, 7, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid.to_grid_mat", "write_shape": [2, 7, 9], "chunk_shape": [2, 7, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_1.expand_index", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_2.expand_index", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pos_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.fc_0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.fc_0.kernel", "write_shape": [256, 128], "chunk_shape": [256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.map.to_m_var", "write_shape": [4, 4], "chunk_shape": [4, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.mlp.layers.4.kernel", "write_shape": [256, 128], "chunk_shape": [256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.pose_gl.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.coefficient_idx", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid.from_grid_mat", "write_shape": [2, 7, 9], "chunk_shape": [2, 7, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid.to_grid_mat", "write_shape": [2, 7, 9], "chunk_shape": [2, 7, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_1.expand_index", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_2.expand_index", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.decoder.rot_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.1.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.0.2.layers.1.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.1.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.attn_layers.1.2.layers.1.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.d_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.gl.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.base_to_contact", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.dof_embedding", "write_shape": [1, 64, 9], "chunk_shape": [1, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.joint_ranges", "write_shape": [2, 2], "chunk_shape": [2, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.joint_transforms", "write_shape": [2, 2], "chunk_shape": [2, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.kinematics_transforms", "write_shape": [2, 7], "chunk_shape": [2, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.0.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.1.base_to_contact", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.1.dof_embedding", "write_shape": [1, 64, 9], "chunk_shape": [1, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.1.joint_ranges", "write_shape": [2, 2], "chunk_shape": [2, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.1.joint_transforms", "write_shape": [2, 2], "chunk_shape": [2, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.1.kinematics_transforms", "write_shape": [2, 7], "chunk_shape": [2, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.1.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.align_to_approach.0", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.align_to_approach.1", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.base_to_contact", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.dof_embedding", "write_shape": [13, 64, 3], "chunk_shape": [13, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.fingertip_idx", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.fingertip_normals", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.init_pregrasp_joint", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.joint_ranges", "write_shape": [4, 2], "chunk_shape": [4, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.joint_transforms", "write_shape": [4, 6], "chunk_shape": [4, 6], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.kinematics_transforms", "write_shape": [4, 7], "chunk_shape": [4, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.local_fingertip_contact_positions", "write_shape": [1, 6, 3], "chunk_shape": [1, 6, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.2.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.align_to_approach.0", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.align_to_approach.1", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.base_to_contact", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.dof_embedding", "write_shape": [17, 64, 3], "chunk_shape": [17, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.fingertip_idx", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.fingertip_normals", "write_shape": [4, 1], "chunk_shape": [4, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.init_pregrasp_joint", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.joint_ranges", "write_shape": [16, 2], "chunk_shape": [16, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.joint_transforms", "write_shape": [16, 2], "chunk_shape": [16, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.kinematics_transforms", "write_shape": [16, 7], "chunk_shape": [16, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.local_fingertip_contact_positions", "write_shape": [4, 1, 3], "chunk_shape": [4, 1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.3.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.4.base_to_contact", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.4.dof_embedding", "write_shape": [23, 64, 3], "chunk_shape": [23, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.4.joint_ranges", "write_shape": [22, 2], "chunk_shape": [22, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.4.joint_transforms", "write_shape": [22, 2], "chunk_shape": [22, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.4.kinematics_transforms", "write_shape": [22, 7], "chunk_shape": [22, 7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.4.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.5.base_to_contact", "write_shape": [7], "chunk_shape": [7], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.5.dof_embedding", "write_shape": [1, 64, 3], "chunk_shape": [1, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.5.joint_ranges", "write_shape": [1, 2], "chunk_shape": [1, 2], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.5.joint_transforms", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.5.kinematics_transforms", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.kinematic_net.kins.5.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_dist.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_dist.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_q_dot_sh.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_q_dot_sh.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.q_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.q_dot_sh_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.kinematic_encoding.weight_irrep.rtoi", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.0.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.0.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.0.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.expand_index", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.1.s_proj.l_weights", "write_shape": [10240], "chunk_shape": [10240], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.layer.l_weights", "write_shape": [20480], "chunk_shape": [20480], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.q_dot_sh_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.q_s_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.2.weight_irrep.rtoi", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_feat.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_feat.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_query.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.ln_query.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.d_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.gl.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.0.4.weight_irrep.rtoi", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.0.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.0.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.0.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.expand_index", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.1.s_proj.l_weights", "write_shape": [10240], "chunk_shape": [10240], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.layer.l_weights", "write_shape": [20480], "chunk_shape": [20480], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.q_dot_sh_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.q_s_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.2.weight_irrep.rtoi", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_feat.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_feat.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_query.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.ln_query.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.d_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.gl.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.multi_scale_feature_query.readout_blocks.1.4.weight_irrep.rtoi", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.attn_weights", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.fnn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn_ln.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ffn_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.kernel", "write_shape": [64, 224], "chunk_shape": [64, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.kernel", "write_shape": [32, 96], "chunk_shape": [32, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 64], "chunk_shape": [128, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.tp_ln.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.tp_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.graph.tp_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_ffn.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_ffn.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_ffn.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_graph.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.0.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.fnn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn_ln.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ffn_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.kernel", "write_shape": [64, 224], "chunk_shape": [64, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.kernel", "write_shape": [32, 96], "chunk_shape": [32, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 64], "chunk_shape": [128, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.tp_ln.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.tp_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.graph.tp_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_ffn.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_ffn.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_ffn.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_graph.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.0.1.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ffn_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.tp_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.graph.tp_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_ffn.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_ffn.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.0.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ffn_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.tp_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.graph.tp_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_ffn.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_ffn.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.1.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ffn_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.tp_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.graph.tp_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_ffn.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_ffn.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.1.2.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ffn_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.tp_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.graph.tp_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_ffn.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_ffn.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.0.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ffn_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.tp_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.graph.tp_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_ffn.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_ffn.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.1.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ffn_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.tp_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.graph.tp_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_ffn.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_ffn.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.mp_blocks.2.2.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.fnn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn_ln.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ffn_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.fc_0.kernel", "write_shape": [64, 224], "chunk_shape": [64, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.fc_0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.fc_0.kernel", "write_shape": [32, 96], "chunk_shape": [32, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 64], "chunk_shape": [128, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.tp_ln.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.tp_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.graph.tp_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_ffn.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_ffn.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_ffn.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_graph.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.proj.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.0.proj.weight", "write_shape": [1, 64, 32], "chunk_shape": [1, 64, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.coefficient_idx", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.grid.from_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.grid.to_grid_mat", "write_shape": [8, 3, 16], "chunk_shape": [8, 3, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_1.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_2.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ffn_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.map.to_m_var", "write_shape": [3, 9], "chunk_shape": [3, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.tp_ln.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.graph.tp_ln.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_ffn.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_ffn.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_graph.balance_degree_weight", "write_shape": [1, 3], "chunk_shape": [1, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.ln_graph.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.proj.expand_index", "write_shape": [3], "chunk_shape": [3], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.pool_blocks.1.proj.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.seed_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.seed_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.seed_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.seed_mlp.layers.2.kernel", "write_shape": [64, 32], "chunk_shape": [64, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_irrep.rtoi", "write_shape": [1, 9], "chunk_shape": [1, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "0.unet.weight_mlp.layers.2.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.0.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.attn_layers.1.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.d_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_gl.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.4.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.dof_mlp.layers.4.kernel", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_dist_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_dist_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_dist_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_dist_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_tmp_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_tmp_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_tmp_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.ln_tmp_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pos_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.fc_0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.fc_0.kernel", "write_shape": [256, 128], "chunk_shape": [256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.mlp.layers.4.kernel", "write_shape": [256, 128], "chunk_shape": [256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.pose_gl.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.decoder.rot_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.1.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.0.2.layers.1.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.1.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.attn_layers.1.2.layers.1.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.d_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.gl.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.0.dof_embedding", "write_shape": [1, 64, 9], "chunk_shape": [1, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.0.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.1.dof_embedding", "write_shape": [1, 64, 9], "chunk_shape": [1, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.1.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.2.dof_embedding", "write_shape": [13, 64, 3], "chunk_shape": [13, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.2.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.3.dof_embedding", "write_shape": [17, 64, 3], "chunk_shape": [17, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.3.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.4.dof_embedding", "write_shape": [23, 64, 3], "chunk_shape": [23, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.4.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.5.dof_embedding", "write_shape": [1, 64, 3], "chunk_shape": [1, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.kinematic_net.kins.5.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_dist.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_dist.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_q_dot_sh.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_q_dot_sh.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.q_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.kinematic_encoding.q_dot_sh_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.0.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.1.s_proj.l_weights", "write_shape": [10240], "chunk_shape": [10240], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.layer.l_weights", "write_shape": [20480], "chunk_shape": [20480], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.q_dot_sh_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.2.q_s_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.d_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.gl.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.0.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.1.s_proj.l_weights", "write_shape": [10240], "chunk_shape": [10240], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.layer.l_weights", "write_shape": [20480], "chunk_shape": [20480], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.q_dot_sh_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.2.q_s_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.d_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.gl.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.attn_weights", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.fnn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ffn_ln.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.kernel", "write_shape": [64, 224], "chunk_shape": [64, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.kernel", "write_shape": [32, 96], "chunk_shape": [32, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 64], "chunk_shape": [128, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.graph.tp_ln.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_ffn.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.0.ln_graph.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.fnn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ffn_ln.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.kernel", "write_shape": [64, 224], "chunk_shape": [64, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.kernel", "write_shape": [32, 96], "chunk_shape": [32, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 64], "chunk_shape": [128, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.graph.tp_ln.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_ffn.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.0.1.ln_graph.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.0.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.1.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.1.2.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.0.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.1.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.mp_blocks.2.2.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.fnn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ffn_ln.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.fc_0.kernel", "write_shape": [64, 224], "chunk_shape": [64, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.fc_0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.fc_0.kernel", "write_shape": [32, 96], "chunk_shape": [32, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 64], "chunk_shape": [128, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.graph.tp_ln.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_ffn.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.ln_graph.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.0.proj.weight", "write_shape": [1, 64, 32], "chunk_shape": [1, 64, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.pool_blocks.1.proj.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.seed_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.seed_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.seed_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.seed_mlp.layers.2.kernel", "write_shape": [64, 32], "chunk_shape": [64, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.weight_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.weight_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.weight_mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.mu.unet.weight_mlp.layers.2.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.0.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.key.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.key.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.out.kernel", "write_shape": [4, 16, 64], "chunk_shape": [4, 16, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.query.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.query.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.value.bias", "write_shape": [4, 16], "chunk_shape": [4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.1.value.kernel", "write_shape": [64, 4, 16], "chunk_shape": [64, 4, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.attn_layers.1.2.layers.1.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.d_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_gl.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.4.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.dof_mlp.layers.4.kernel", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_dist_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_dist_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_dist_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_dist_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_tmp_dof.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_tmp_dof.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_tmp_pose.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.ln_tmp_pose.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pos_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.fc_0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.fc_0.kernel", "write_shape": [256, 128], "chunk_shape": [256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.mlp.layers.4.kernel", "write_shape": [256, 128], "chunk_shape": [256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.pose_gl.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.grid_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.grid_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.scalar_mlp.layers.4.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.so3_linear_1.weight", "write_shape": [2, 64, 64], "chunk_shape": [2, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.so3_linear_2.bias", "write_shape": [1], "chunk_shape": [1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.decoder.rot_head.so3_linear_2.weight", "write_shape": [2, 1, 64], "chunk_shape": [2, 1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.1.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.0.2.layers.1.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.1.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.2.layers.1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.attn_layers.1.2.layers.1.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.d_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.gl.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.0.dof_embedding", "write_shape": [1, 64, 9], "chunk_shape": [1, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.0.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.1.dof_embedding", "write_shape": [1, 64, 9], "chunk_shape": [1, 64, 9], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.1.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.2.dof_embedding", "write_shape": [13, 64, 3], "chunk_shape": [13, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.2.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.3.dof_embedding", "write_shape": [17, 64, 3], "chunk_shape": [17, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.3.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.4.dof_embedding", "write_shape": [23, 64, 3], "chunk_shape": [23, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.4.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.5.dof_embedding", "write_shape": [1, 64, 3], "chunk_shape": [1, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.kinematic_net.kins.5.pose_embedding", "write_shape": [2, 64, 3], "chunk_shape": [2, 64, 3], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_dist.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_dist.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_q_dot_sh.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_q_dot_sh.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.q_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.kinematic_encoding.q_dot_sh_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.0.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.1.s_proj.l_weights", "write_shape": [10240], "chunk_shape": [10240], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.layer.l_weights", "write_shape": [20480], "chunk_shape": [20480], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.mlp.layers.3.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.q_dot_sh_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.2.q_s_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.ln_feat.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.ln_query.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.0.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.attn_layers.1.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.d_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.gl.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.0.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.0.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.0.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_scalar_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.film.time_vec_scales.layers.4.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.0.1.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.1.1.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.per_scale_ln.2.1.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.1.s_proj.l_weights", "write_shape": [10240], "chunk_shape": [10240], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.layer.d_weights", "write_shape": [1216], "chunk_shape": [1216], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.layer.l_weights", "write_shape": [20480], "chunk_shape": [20480], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_q_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_q_feat.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_scalar.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_time.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.ln_time.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.key.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.key.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.out.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.out.kernel", "write_shape": [8, 32, 256], "chunk_shape": [8, 32, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.query.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.query.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.value.bias", "write_shape": [8, 32], "chunk_shape": [8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mha.value.kernel", "write_shape": [256, 8, 32], "chunk_shape": [256, 8, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.1.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.mlp.layers.3.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.q_dot_sh_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.2.q_s_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.ln_feat.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.ln_query.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.biases.bias_0", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.0", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.1", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.2", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.3", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.4", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.5", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.6", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.7", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.8", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.9", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.10", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.11", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.12", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.13", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.3.tp.weights.14", "write_shape": [64, 64, 64], "chunk_shape": [64, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.0.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.key.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.out.kernel", "write_shape": [2, 32, 192], "chunk_shape": [2, 32, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.query.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.bias", "write_shape": [2, 32], "chunk_shape": [2, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.attn_layers.1.value.kernel", "write_shape": [64, 6, 32], "chunk_shape": [64, 6, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_mlp.layers.4.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.attn_weights", "write_shape": [1, 64], "chunk_shape": [1, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.d_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.ln_sd_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.0.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.2.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.4.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.msg_mlp.layers.6.kernel", "write_shape": [256, 64], "chunk_shape": [256, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.gl.sd_dot_modulation", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_dist_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.multi_scale_feature_query.readout_blocks.1.4.ln_time_graph.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.attn_weights", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.fnn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ffn_ln.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.fc_0.kernel", "write_shape": [64, 224], "chunk_shape": [64, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.fc_0.kernel", "write_shape": [32, 96], "chunk_shape": [32, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 64], "chunk_shape": [128, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.graph.tp_ln.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_ffn.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.0.ln_graph.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.fnn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ffn_ln.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.fc_0.kernel", "write_shape": [64, 224], "chunk_shape": [64, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.fc_0.kernel", "write_shape": [32, 96], "chunk_shape": [32, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 64], "chunk_shape": [128, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.graph.tp_ln.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_ffn.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.0.1.ln_graph.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.0.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.1.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.1.2.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.0.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.1.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.mp_blocks.2.2.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.fnn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.attn_weights", "write_shape": [4, 32], "chunk_shape": [4, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [32, 32], "chunk_shape": [32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.so3_linear_1.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.so3_linear_1.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.so3_linear_2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn.so3_linear_2.weight", "write_shape": [1, 32, 32], "chunk_shape": [1, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ffn_ln.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ln_d_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ln_d_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ln_s_dot.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.ln_s_dot.scale", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.fc_0.bias", "write_shape": [224], "chunk_shape": [224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.fc_0.kernel", "write_shape": [64, 224], "chunk_shape": [64, 224], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.fc_0.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.fc_0.kernel", "write_shape": [32, 96], "chunk_shape": [32, 96], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [128, 64], "chunk_shape": [128, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [32, 64], "chunk_shape": [32, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.tp_ln.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.graph.tp_ln.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_ffn.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_ffn.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_graph.affine_bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.ln_graph.affine_weight", "write_shape": [32, 1], "chunk_shape": [32, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.0.proj.weight", "write_shape": [1, 64, 32], "chunk_shape": [1, 64, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.fnn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.attn_weights", "write_shape": [4, 64], "chunk_shape": [4, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.grid_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.scalar_mlp.layers.2.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.so3_linear_1.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.so3_linear_1.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.so3_linear_2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn.so3_linear_2.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ffn_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ln_d_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ln_d_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ln_s_dot.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.ln_s_dot.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.fc_0.bias", "write_shape": [448], "chunk_shape": [448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.fc_0.kernel", "write_shape": [128, 448], "chunk_shape": [128, 448], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.bias", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.mlp.layers.4.kernel", "write_shape": [64, 768], "chunk_shape": [64, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.0.fc.kernel", "write_shape": [256, 256], "chunk_shape": [256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp1.so2_m_conv.1.fc.kernel", "write_shape": [128, 128], "chunk_shape": [128, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.fc_0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.fc_0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.0.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.2.kernel", "write_shape": [64, 192], "chunk_shape": [64, 192], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.mlp.layers.4.kernel", "write_shape": [64, 384], "chunk_shape": [64, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.0.fc.kernel", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.so2_tp2.so2_m_conv.1.fc.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.tp_ln.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.graph.tp_ln.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_attr.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_attr.0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_ffn.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_ffn.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_graph.affine_bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.ln_graph.affine_weight", "write_shape": [64, 1], "chunk_shape": [64, 1], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.proj.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.pool_blocks.1.proj.weight", "write_shape": [1, 64, 64], "chunk_shape": [1, 64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.seed_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.seed_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.seed_mlp.layers.2.bias", "write_shape": [32], "chunk_shape": [32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.seed_mlp.layers.2.kernel", "write_shape": [64, 32], "chunk_shape": [64, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.weight_mlp.layers.0.bias", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.weight_mlp.layers.0.kernel", "write_shape": [64, 64], "chunk_shape": [64, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.weight_mlp.layers.2.bias", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.0.nu.unet.weight_mlp.layers.2.kernel", "write_shape": [64, 128], "chunk_shape": [64, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.opt_state.2.count", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}, {"array_metadata": {"param_name": "1.step", "write_shape": [], "chunk_shape": [], "ext_metadata": null}}]}
\ No newline at end of file
diff --git a/me-seed-anneal_25000_4/d/6df16f8850d87665698b0829137fd441 b/me-seed-anneal_25000_4/d/6df16f8850d87665698b0829137fd441
new file mode 100644
index 0000000000000000000000000000000000000000..722ff4bd57740d802cb4db306799a8ee9b9ceb2c
--- /dev/null
+++ b/me-seed-anneal_25000_4/d/6df16f8850d87665698b0829137fd441
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a7b643fa90aee479812581bd10e849f6e448e0a07de3dcb08cf1f238e78508a8
+size 837407
diff --git a/me-seed-anneal_25000_4/manifest.ocdbt b/me-seed-anneal_25000_4/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..5212a24c2ef892e38f27ac1a5ba31719555b4fc1
Binary files /dev/null and b/me-seed-anneal_25000_4/manifest.ocdbt differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/03ad492c8c39ffec2da8487562b7be84 b/me-seed-anneal_25000_4/ocdbt.process_0/d/03ad492c8c39ffec2da8487562b7be84
new file mode 100644
index 0000000000000000000000000000000000000000..173db2c7366347f570f52438918a1a6f0ffb9e86
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/03ad492c8c39ffec2da8487562b7be84 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/055a4101090e22d0f730e9c4ababccbc b/me-seed-anneal_25000_4/ocdbt.process_0/d/055a4101090e22d0f730e9c4ababccbc
new file mode 100644
index 0000000000000000000000000000000000000000..d51d893bf5fbc1667f8a13f6a1ed67a8e5af6e27
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/055a4101090e22d0f730e9c4ababccbc differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/0db36d38697138532b242912e2c20495 b/me-seed-anneal_25000_4/ocdbt.process_0/d/0db36d38697138532b242912e2c20495
new file mode 100644
index 0000000000000000000000000000000000000000..11a037bed5d196e9b60d3fc43b6533088d09012e
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/0db36d38697138532b242912e2c20495 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/15246916a0fbe858dbebab383f62aa8e b/me-seed-anneal_25000_4/ocdbt.process_0/d/15246916a0fbe858dbebab383f62aa8e
new file mode 100644
index 0000000000000000000000000000000000000000..4f7e67c67fa42118e57808e21ebb0135dfd056f8
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/15246916a0fbe858dbebab383f62aa8e differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/1941934e90478a89b6483c9ac85293c5 b/me-seed-anneal_25000_4/ocdbt.process_0/d/1941934e90478a89b6483c9ac85293c5
new file mode 100644
index 0000000000000000000000000000000000000000..4672312e99da37c1418a27507d5db84313304784
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/1941934e90478a89b6483c9ac85293c5 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/289eca1070fe367db622017eca77ccc9 b/me-seed-anneal_25000_4/ocdbt.process_0/d/289eca1070fe367db622017eca77ccc9
new file mode 100644
index 0000000000000000000000000000000000000000..005b2d0874cf37e1f22e6056bd7ffc0685464984
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/289eca1070fe367db622017eca77ccc9 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/2f386fcb89d33f57e980f9c648b50af4 b/me-seed-anneal_25000_4/ocdbt.process_0/d/2f386fcb89d33f57e980f9c648b50af4
new file mode 100644
index 0000000000000000000000000000000000000000..19ece18c7693e87d91fb52fbc809276ae0d5b1ad
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/2f386fcb89d33f57e980f9c648b50af4 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/3274bfb4d77bcef18e0a460dd0d3be2c b/me-seed-anneal_25000_4/ocdbt.process_0/d/3274bfb4d77bcef18e0a460dd0d3be2c
new file mode 100644
index 0000000000000000000000000000000000000000..559847b53d4bab9da223d71c912210bc9c65b91e
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/3274bfb4d77bcef18e0a460dd0d3be2c differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/3846cd84260ac91bb6af2590e13b8f8e b/me-seed-anneal_25000_4/ocdbt.process_0/d/3846cd84260ac91bb6af2590e13b8f8e
new file mode 100644
index 0000000000000000000000000000000000000000..cd6c395a62b265c957b5a4104f589a4b543b9a5f
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/3846cd84260ac91bb6af2590e13b8f8e differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/391943bcd94856bc562d6f6d9959d7e5 b/me-seed-anneal_25000_4/ocdbt.process_0/d/391943bcd94856bc562d6f6d9959d7e5
new file mode 100644
index 0000000000000000000000000000000000000000..f68c9dd79e22cd47ec1e39792fce0f24c706f543
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/391943bcd94856bc562d6f6d9959d7e5 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/3d0f0c9004bd81f01d77a5f50e8e3329 b/me-seed-anneal_25000_4/ocdbt.process_0/d/3d0f0c9004bd81f01d77a5f50e8e3329
new file mode 100644
index 0000000000000000000000000000000000000000..9e8c4608fbaefb6feba60545bbde3ca3a33456cd
--- /dev/null
+++ b/me-seed-anneal_25000_4/ocdbt.process_0/d/3d0f0c9004bd81f01d77a5f50e8e3329
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:85505832501613f8ec206692218bcc91708ec871783debbed5a2e7613636204d
+size 6868992
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/438c729771711f7537775abf0c7c3cd8 b/me-seed-anneal_25000_4/ocdbt.process_0/d/438c729771711f7537775abf0c7c3cd8
new file mode 100644
index 0000000000000000000000000000000000000000..7a7c2449d7348d6695d8943838e872bbdfb94063
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/438c729771711f7537775abf0c7c3cd8 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/4ace837bfb2fbb91f6b7b4c3b9f9f6ea b/me-seed-anneal_25000_4/ocdbt.process_0/d/4ace837bfb2fbb91f6b7b4c3b9f9f6ea
new file mode 100644
index 0000000000000000000000000000000000000000..c5c3b9a6f80d16b4bdc09fc2fdb6f94c1eb85400
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/4ace837bfb2fbb91f6b7b4c3b9f9f6ea differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/4bd55f6519b201437dd556f54b268fad b/me-seed-anneal_25000_4/ocdbt.process_0/d/4bd55f6519b201437dd556f54b268fad
new file mode 100644
index 0000000000000000000000000000000000000000..c276950e57c18a76985de2f7d8b9594a8facf51d
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/4bd55f6519b201437dd556f54b268fad differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/4be70cfae974d68b52713e78cd01e485 b/me-seed-anneal_25000_4/ocdbt.process_0/d/4be70cfae974d68b52713e78cd01e485
new file mode 100644
index 0000000000000000000000000000000000000000..dfdced4758ab1ae1830516b32dc1cbc2e639ce59
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/4be70cfae974d68b52713e78cd01e485 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/52a4586ccd77904f4711c95526687f22 b/me-seed-anneal_25000_4/ocdbt.process_0/d/52a4586ccd77904f4711c95526687f22
new file mode 100644
index 0000000000000000000000000000000000000000..f13956e88fb5c801c6856db42489f93518490612
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/52a4586ccd77904f4711c95526687f22 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/56526694f7a4e49862307a97ebc34913 b/me-seed-anneal_25000_4/ocdbt.process_0/d/56526694f7a4e49862307a97ebc34913
new file mode 100644
index 0000000000000000000000000000000000000000..e789166f085271b33958e54a2df24f87eccb9c1b
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/56526694f7a4e49862307a97ebc34913 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/60216d6ce7a9a1bc5e6a58d389d57f94 b/me-seed-anneal_25000_4/ocdbt.process_0/d/60216d6ce7a9a1bc5e6a58d389d57f94
new file mode 100644
index 0000000000000000000000000000000000000000..4478da1967ebfa2f1e822a76c51c26ed88be1e43
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/60216d6ce7a9a1bc5e6a58d389d57f94 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/60aca8647e3d5929295e5ceaf9c045c3 b/me-seed-anneal_25000_4/ocdbt.process_0/d/60aca8647e3d5929295e5ceaf9c045c3
new file mode 100644
index 0000000000000000000000000000000000000000..10bb934b26e8c8deb298f4292cffcef902f0d648
--- /dev/null
+++ b/me-seed-anneal_25000_4/ocdbt.process_0/d/60aca8647e3d5929295e5ceaf9c045c3
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3af0eafa771d9b16a90e04e937575ebbfb47896521daba0fdcb5964a8168bedc
+size 45211648
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/61b832ad8f494c0d592317db65395e6b b/me-seed-anneal_25000_4/ocdbt.process_0/d/61b832ad8f494c0d592317db65395e6b
new file mode 100644
index 0000000000000000000000000000000000000000..bb5364477b3001ebcd04eb92acf02af823c18768
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/61b832ad8f494c0d592317db65395e6b differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/662f3479fd7e24fc8ab9983df2de613d b/me-seed-anneal_25000_4/ocdbt.process_0/d/662f3479fd7e24fc8ab9983df2de613d
new file mode 100644
index 0000000000000000000000000000000000000000..be5b727c6633c83e21db88a1ad1152bd13acb32f
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/662f3479fd7e24fc8ab9983df2de613d differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/66e8ab822b8731a51a7df8a7a5838836 b/me-seed-anneal_25000_4/ocdbt.process_0/d/66e8ab822b8731a51a7df8a7a5838836
new file mode 100644
index 0000000000000000000000000000000000000000..9e4881d439f1e4d122a084cf5c08dfc48131a470
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/66e8ab822b8731a51a7df8a7a5838836 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/77606384613d5a6b089c5753859c8441 b/me-seed-anneal_25000_4/ocdbt.process_0/d/77606384613d5a6b089c5753859c8441
new file mode 100644
index 0000000000000000000000000000000000000000..25beba543fd8b8c4b53353c5cb121a135f47901e
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/77606384613d5a6b089c5753859c8441 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/777dd520282305a722e9a0640c2d0805 b/me-seed-anneal_25000_4/ocdbt.process_0/d/777dd520282305a722e9a0640c2d0805
new file mode 100644
index 0000000000000000000000000000000000000000..8bfb332eab7d9bbddbf9c607d31def024d282dae
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/777dd520282305a722e9a0640c2d0805 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/7d73d15c978a73c416e797f77148e96b b/me-seed-anneal_25000_4/ocdbt.process_0/d/7d73d15c978a73c416e797f77148e96b
new file mode 100644
index 0000000000000000000000000000000000000000..5ce47142b1e64c8df607808ae26ecad0ce433cca
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/7d73d15c978a73c416e797f77148e96b differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/89bad5ebf2fc55101c07f53d0ad3137b b/me-seed-anneal_25000_4/ocdbt.process_0/d/89bad5ebf2fc55101c07f53d0ad3137b
new file mode 100644
index 0000000000000000000000000000000000000000..7492b93e9279744ce73d6cca3e46e4607003a175
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/89bad5ebf2fc55101c07f53d0ad3137b differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/919302d03986202605d9e9b782d67153 b/me-seed-anneal_25000_4/ocdbt.process_0/d/919302d03986202605d9e9b782d67153
new file mode 100644
index 0000000000000000000000000000000000000000..3c8e3966b140036961d4899c0cd610f0e901feec
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/919302d03986202605d9e9b782d67153 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/9d01d4f4ee8ed87496459df922511e4f b/me-seed-anneal_25000_4/ocdbt.process_0/d/9d01d4f4ee8ed87496459df922511e4f
new file mode 100644
index 0000000000000000000000000000000000000000..496cceb4e8c1929839872692fea856ab673f1c42
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/9d01d4f4ee8ed87496459df922511e4f differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/9d9e4e0a390292d22d337dfe778cfde4 b/me-seed-anneal_25000_4/ocdbt.process_0/d/9d9e4e0a390292d22d337dfe778cfde4
new file mode 100644
index 0000000000000000000000000000000000000000..86fba4976bdf149a78a60b2b0f8d16085f458050
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/9d9e4e0a390292d22d337dfe778cfde4 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/a0b5d0bb7f2ad0b02d8a64669ff910b8 b/me-seed-anneal_25000_4/ocdbt.process_0/d/a0b5d0bb7f2ad0b02d8a64669ff910b8
new file mode 100644
index 0000000000000000000000000000000000000000..01630cc55dd5c20bc544b6b66bb759a0c3a30a0b
--- /dev/null
+++ b/me-seed-anneal_25000_4/ocdbt.process_0/d/a0b5d0bb7f2ad0b02d8a64669ff910b8
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:31030251131ad2eb0841bf9933e6c846d043bad282ee682c838204755c9f55fd
+size 61366272
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/a141ec9430e4b8e5cfa5150687c04a32 b/me-seed-anneal_25000_4/ocdbt.process_0/d/a141ec9430e4b8e5cfa5150687c04a32
new file mode 100644
index 0000000000000000000000000000000000000000..83d90cf49890129d8f6b765a9207933206f318c1
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/a141ec9430e4b8e5cfa5150687c04a32 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/af60d77cbc143d1bd4bb1e80cbbbbe5d b/me-seed-anneal_25000_4/ocdbt.process_0/d/af60d77cbc143d1bd4bb1e80cbbbbe5d
new file mode 100644
index 0000000000000000000000000000000000000000..43e0d728a62cc62d8d03dc7e30da63a36a750c05
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/af60d77cbc143d1bd4bb1e80cbbbbe5d differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/b3210342df51754fa334e2473712e5d8 b/me-seed-anneal_25000_4/ocdbt.process_0/d/b3210342df51754fa334e2473712e5d8
new file mode 100644
index 0000000000000000000000000000000000000000..815f520c68b603018ddaabc6a3be75c1ab8faf9e
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/b3210342df51754fa334e2473712e5d8 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/b8b89906351e11b64fc3eeb9a44e8408 b/me-seed-anneal_25000_4/ocdbt.process_0/d/b8b89906351e11b64fc3eeb9a44e8408
new file mode 100644
index 0000000000000000000000000000000000000000..c48c4c7e5ff4792239550bcabb0e887a6a0272cf
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/b8b89906351e11b64fc3eeb9a44e8408 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/babf834dbeabbdd6b108113607f72684 b/me-seed-anneal_25000_4/ocdbt.process_0/d/babf834dbeabbdd6b108113607f72684
new file mode 100644
index 0000000000000000000000000000000000000000..1103f263a5e9b4a38e764dfd58d4cf6a22246761
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/babf834dbeabbdd6b108113607f72684 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/c03bc91493126442b153a40c8cc82b9b b/me-seed-anneal_25000_4/ocdbt.process_0/d/c03bc91493126442b153a40c8cc82b9b
new file mode 100644
index 0000000000000000000000000000000000000000..ce5a69c0a7361bbfa638b25bb88cfe5e5e03336a
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/c03bc91493126442b153a40c8cc82b9b differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/c581081f89a801613782f7d6db4262b2 b/me-seed-anneal_25000_4/ocdbt.process_0/d/c581081f89a801613782f7d6db4262b2
new file mode 100644
index 0000000000000000000000000000000000000000..164b69418accff787c3fa0d2cb37ca36e9e07da6
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/c581081f89a801613782f7d6db4262b2 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/c8eaad476b04d53e07ab7bc39e9cade2 b/me-seed-anneal_25000_4/ocdbt.process_0/d/c8eaad476b04d53e07ab7bc39e9cade2
new file mode 100644
index 0000000000000000000000000000000000000000..94de20826ea5e6e911a43b6fee4578912af43ae6
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/c8eaad476b04d53e07ab7bc39e9cade2 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/c8eb449374e3ee424492a8d3b3d0faa9 b/me-seed-anneal_25000_4/ocdbt.process_0/d/c8eb449374e3ee424492a8d3b3d0faa9
new file mode 100644
index 0000000000000000000000000000000000000000..cb14060f44607274e7d940c821dffc07747644b8
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/c8eb449374e3ee424492a8d3b3d0faa9 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/d116653203ff60d44da6f4a13b1cf50c b/me-seed-anneal_25000_4/ocdbt.process_0/d/d116653203ff60d44da6f4a13b1cf50c
new file mode 100644
index 0000000000000000000000000000000000000000..ea97d2c8f93e54a987e369a69287f6e4d1db9080
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/d116653203ff60d44da6f4a13b1cf50c differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/d67641bd3d34cbd39e0ab8104b567e2d b/me-seed-anneal_25000_4/ocdbt.process_0/d/d67641bd3d34cbd39e0ab8104b567e2d
new file mode 100644
index 0000000000000000000000000000000000000000..10233e0701eaa12bae0e1e62b09092676cfed87e
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/d67641bd3d34cbd39e0ab8104b567e2d differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/e11fb568cc86faf16b7f54796c23ca83 b/me-seed-anneal_25000_4/ocdbt.process_0/d/e11fb568cc86faf16b7f54796c23ca83
new file mode 100644
index 0000000000000000000000000000000000000000..32de991f86000b86eb50844e0dac63860fd64328
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/e11fb568cc86faf16b7f54796c23ca83 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/e6f4a07d5a6435f3c5dab930d72e9ae6 b/me-seed-anneal_25000_4/ocdbt.process_0/d/e6f4a07d5a6435f3c5dab930d72e9ae6
new file mode 100644
index 0000000000000000000000000000000000000000..e06a99f852f6938f3aff313c6b9aa30990f668ab
--- /dev/null
+++ b/me-seed-anneal_25000_4/ocdbt.process_0/d/e6f4a07d5a6435f3c5dab930d72e9ae6
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:893f1024d787cfd64cdfcf398fa4c1f023ce5eb05dcfd1323a3934a42ccdee93
+size 93970432
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/ee21af9be2478482ae76f5eed556031d b/me-seed-anneal_25000_4/ocdbt.process_0/d/ee21af9be2478482ae76f5eed556031d
new file mode 100644
index 0000000000000000000000000000000000000000..63dac40ec63ad1915cc19f4c5e7ad382d869beea
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/ee21af9be2478482ae76f5eed556031d differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/effa339b7614f4cb04f5859d0a52de78 b/me-seed-anneal_25000_4/ocdbt.process_0/d/effa339b7614f4cb04f5859d0a52de78
new file mode 100644
index 0000000000000000000000000000000000000000..f6c858b14ceb449b4b67512d6a8634148377b812
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/effa339b7614f4cb04f5859d0a52de78 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/f3fdcd4704bdab6d2a1dd8c53d9859f1 b/me-seed-anneal_25000_4/ocdbt.process_0/d/f3fdcd4704bdab6d2a1dd8c53d9859f1
new file mode 100644
index 0000000000000000000000000000000000000000..40ed873f94820d306a960eb009172582948390e5
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/f3fdcd4704bdab6d2a1dd8c53d9859f1 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/f42fb4ba34078c0b98e09e31fe923cf0 b/me-seed-anneal_25000_4/ocdbt.process_0/d/f42fb4ba34078c0b98e09e31fe923cf0
new file mode 100644
index 0000000000000000000000000000000000000000..415168ff357b9ef741c6935ebc612fe92a92065c
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/f42fb4ba34078c0b98e09e31fe923cf0 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/d/fc4eeddf4d02027316f8e57e86682c59 b/me-seed-anneal_25000_4/ocdbt.process_0/d/fc4eeddf4d02027316f8e57e86682c59
new file mode 100644
index 0000000000000000000000000000000000000000..9331ee20e7122e643bdbfabf4f39d05211f0cb41
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/d/fc4eeddf4d02027316f8e57e86682c59 differ
diff --git a/me-seed-anneal_25000_4/ocdbt.process_0/manifest.ocdbt b/me-seed-anneal_25000_4/ocdbt.process_0/manifest.ocdbt
new file mode 100644
index 0000000000000000000000000000000000000000..d4f246fdfda179205bf40bd9b57ad4ff20bb60c9
Binary files /dev/null and b/me-seed-anneal_25000_4/ocdbt.process_0/manifest.ocdbt differ